Reduce Memory Use

Use the memory APIs to size a run, select a supported checkpointing option, and identify the largest static memory costs. The limits of each report are listed with the relevant API below.

For S-parameters, ports, and far-field results, the supported solver and calculator combinations remain those documented in the relevant support guide. Changing the mesh or checkpointing configuration does not extend that support.

Current estimate_ad_memory(...), plan_ad_memory(...), and explain_ad_memory(...) artifacts are not certificates. Their evidence_class values distinguish the type of information returned:

Run ad_memory_preflight(...) before a memory-heavy AD or inverse-design job. It calls plan_ad_memory(...), explain_ad_memory(...), mesh_intelligence_report(...), and, when requested, diagnose_ad_saved_residuals(...). It does not run FDTD.

Setup used by the snippets below:

import jax
import jax.numpy as jnp
from rfx import Box, GaussianPulse, Simulation

sim = Simulation(
    freq_max=5e9,
    domain=(0.02, 0.02, 0.01),
    dx=2e-3,
    boundary="cpml",
    cpml_layers=4,
)
sim.add_material("slab", eps_r=4.0, sigma=0.01)
sim.add(Box((0.006, 0.006, 0.003), (0.014, 0.014, 0.007)), material="slab")
sim.add_source(
    (0.005, 0.010, 0.005),
    "ez",
    waveform=GaussianPulse(f0=3e9, bandwidth=0.8),
    amplitude_kind="current",
)
sim.add_probe((0.015, 0.010, 0.005), "ez")

# Names the later snippets read: the design field and the region it is
# allowed to vary in, and a stand-in for whatever launches the real run.
eps = jnp.full(sim.run(n_steps=100).grid.shape, 4.0, dtype=jnp.float32)
region_mask = jnp.zeros_like(eps, dtype=bool).at[4:-4, 4:-4, 4:-4].set(True)


def launch():
    print("start the configured run here")

preflight = sim.ad_memory_preflight(
    n_steps=10_000,
    available_memory_gb=24.0,
)

if preflight.full_ad_fits:
    checkpoint_kwargs = {}
elif (
    preflight.checkpointing_fits
    and preflight.supported_checkpoint_mode == "checkpoint_every"
):
    checkpoint_kwargs = {"checkpoint_every": preflight.checkpoint_every}
elif (
    preflight.checkpointing_fits
    and preflight.supported_checkpoint_mode == "checkpoint_segments"
):
    checkpoint_kwargs = {"checkpoint_segments": preflight.checkpoint_segments}
else:
    # The returned least-memory checkpoint candidate is diagnostic-only here.
    raise RuntimeError(preflight.recommendation)

print(preflight.status)
print(preflight.recommendation)
print([hint.to_dict() for hint in preflight.action_hints])
preflight_json = preflight.to_json()

# After choosing a supported configuration and running the required physics checks,
# pass the selected knob into the launch path:
# result = sim.forward(n_steps=10_000, **checkpoint_kwargs)

To include JAX trace information, pass a small representative objective through residual_fun. residual_context is copied before the diagnostic runs. It accepts JSON primitives, mappings, sequences, NumPy/JAX scalar or array .tolist() values, and valid to_dict() or to_json() outputs. Unsupported objects raise TypeError; non-finite JSON values raise ValueError.

import jax.numpy as jnp

def reduced_loss(params):
    return jnp.sum(jnp.sin(params) ** 2)

preflight = sim.ad_memory_preflight(
    n_steps=10_000,
    available_memory_gb=24.0,
    residual_fun=reduced_loss,
    residual_args=(jnp.ones((32,)),),
    residual_workflow="reduced-inverse-design-loss",
    residual_context={"case": "notch-filter-demo"},
)

The preflight report has these evidence boundaries:

• static memory planning only
• trace-time JAX saved-residual explainability only when residual_diagnostic is present
• not a runtime peak-memory guarantee
• not XLA memory analysis
• not profiler evidence
• not a certificate
• not RF validation

After selecting the configuration, compile the exact AD objective or runner to be launched. Pass that compiled object to ad_memory_compiled_certificate(...), using the same checkpoint arguments and scope as the preflight report. The API reads Compiled.memory_analysis() once and fails closed if the scope or compiler byte fields are incomplete.

compiled_loss = jax.jit(reduced_loss).lower(jnp.ones((32,))).compile()

certificate = sim.ad_memory_compiled_certificate(
    compiled_loss,
    n_steps=10_000,
    available_memory_gb=24.0,
    target_fraction=0.85,
    **checkpoint_kwargs,
    precision="float32",
    input_signature={
        "eps_r": {"shape": [128, 64, 16], "dtype": "float32", "role": "design"}
    },
    static_signature={"dx": 0.0005, "boundary": "pml"},
    compiled_object_id="notch-filter-loss-v1",
    runner_or_objective="notch-filter-reduced-loss",
    preflight=preflight,
)

Signature dictionaries may contain representative NumPy/JAX arrays or jax.ShapeDtypeStruct leaves. rfx stores shape, dtype, and weak-type metadata for those array-like leaves, not full array values. Use explicit dictionary fields such as "role": "design" when a semantic label matters for an audit.

# `compiler_estimate_within_budget` means the JAX compiler's memory estimate
# fits the target budget for this exact scope. Allocator fragmentation and
# runtime scratch are not modeled, so a fit at high utilization can still OOM.
if certificate.status == "compiler_estimate_within_budget":
    launch()
else:
    raise RuntimeError(certificate.recommendations[0])

The certificate scope includes the compiled object, backend/device metadata, input and static signatures, precision, n_steps, warmup count, checkpoint mode, and target budget. The other statuses are:

• compiler_estimate_exceeds_budget: the complete-scope compiler estimate is larger than available_memory_gb * target_fraction; do not launch that scope under the stated budget.
• scope_incomplete: required exact-scope metadata is missing or not JSON-safe.
• scope_mismatch: preflight metadata contradicts the supplied scope or compiled-object introspection.
• analysis_unavailable: memory_analysis() is missing, raises, or returns None.
• analysis_incomplete: a required byte field is missing or invalid.

The required byte fields are temp_size_in_bytes, argument_size_in_bytes, output_size_in_bytes, and alias_size_in_bytes. rfx computes required bytes as temp + argument + output - alias, then compares the result with available_memory_gb * target_fraction.

scope_digest, config_digest, and environment_digest are audit identities over canonical JSON. They do not establish that Python source, callables, or a simulation configuration correspond to an opaque compiled executable unless separate introspection establishes that relationship. JAX memory analysis is estimated, version/backend dependent and may be unavailable. A bounded_certificate is therefore a compiler estimate for one exact scope, not profiler evidence, RF validation, or a runtime peak-memory guarantee.

A non-uniform z mesh can reduce the number of cells when a small feature needs fine resolution along one axis. Use the non-uniform mesh guide to check supported combinations and mesh restrictions.

plan = sim.plan_ad_memory(n_steps=10_000, available_memory_gb=24.0)
if not (plan.full_ad_fits or plan.segmented_fits):
    raise RuntimeError(plan.recommendation)

report = sim.mesh_intelligence_report(
    n_steps=10_000,
    checkpoint_every=plan.checkpoint_every if plan.segmented_fits else None,
    checkpoint_segments=plan.checkpoint_segments if plan.segmented_fits else None,
    available_memory_gb=24.0,
)

print(plan.recommendation)
print(report.cell_savings_factor)
print(report.recommendation)

plan_json = plan.to_json()
report_json = report.to_json()

Use explain_ad_memory(...) to inspect the selected estimate:

explain = sim.explain_ad_memory(
    n_steps=10_000,
    checkpoint_every=plan.checkpoint_every if plan.segmented_fits else None,
    checkpoint_segments=plan.checkpoint_segments if plan.segmented_fits else None,
    available_memory_gb=24.0,
)

print(explain.dominant_component)
print(explain.recommendations)
explain_json = explain.to_json()

The report selects ad_full_gb or ad_segmented_gb and breaks the estimate into full field tape, segmented boundary state, live-segment rematerialization tape, CPML/material state, and NTFF monitor state. The mesh report also returns:

• cell_savings_factor: comparison with a uniform grid at the smallest configured cell size;
• preflight_issues: geometry, resolution, and feature-support warnings;
• ad_memory.ad_segmented_gb and ad_memory.ad_segmented_active_segments: the segmented estimate and active segment count when either segmented checkpoint mode is used;
• recommendation: the suggested next action for the configuration.

plan_ad_memory(...) returns both checkpoint values as None when full AD fits. In that case, full_ad_fits=True and segmented_fits=False because no segmented candidate is needed. Otherwise it selects checkpoint_every on a non-uniform grid or checkpoint_segments on a uniform grid. Use checkpoint_mode only when segmented_fits is true. If both fit flags are false, follow plan.recommendation rather than launching the returned least-memory candidate.

checkpoint_every is a chunk length in timesteps for the non-uniform scan-of-scan implementation. checkpoint_segments is the number of equal segments for the uniform segmented scan and must divide n_steps. The segmented estimate includes carry and cotangent state at active segment boundaries (2 x active_segments field states) plus one live segment’s rematerialization tape. These terms balance near sqrt(2 x n_steps) timesteps per segment; either extreme can increase peak memory, so use the value chosen by plan_ad_memory(...) rather than minimizing the segment count manually.

n_warmup reduces the active reverse-mode timestep count and must satisfy 0 <= n_warmup < n_steps. It is an APPROXIMATION, not merely a memory optimization — see forward()’s own n_warmup docstring for the measured gradient-truncation error curve — and estimate_ad_memory/plan_ad_memory report it as a hypothetical planning number regardless of mesh type; on the uniform (single-device) mesh, forward(n_warmup=...) itself raises NotImplementedError (non-uniform / distributed-non-uniform meshes only). The plan applies a conservative safety multiplier before setting the fit flags and records it in plan.fit_safety_factor.

forward(design_mask=...) was removed entirely (issue #625) rather than deprecated: it was measured to save zero reverse-mode AD memory (JAX’s partial-eval residuals have whole-array granularity, so confining a design variable to a fraction of the grid does not shrink the tape) while corrupting the gradient in every configuration tested. It is not a memory option — do not look for a direct replacement kwarg.

If you actually need to restrict which cells carry a derivative (not for memory — for design-region hygiene, e.g. keeping a gradient from leaking into a PML/absorber region), construct the restriction yourself at the eps_override call site, one line:

eps = jnp.where(region_mask, eps, jax.lax.stop_gradient(eps))
result = sim.forward(eps_override=eps, n_steps=200)

region_mask is a boolean array (or broadcastable to eps’s shape) that is True inside the design region. This is forward-identity (stop_gradient does not change the value, only the gradient), and the resulting per-cell gradient is exactly zero outside region_mask and exactly equal to the unmasked gradient inside it. For actual memory reduction, use checkpoint_every / checkpoint_segments above instead.

For a uniform run, gate execution on the fit flags before applying the selected checkpoint value:

plan = sim.plan_ad_memory(n_steps=10_000, available_memory_gb=24.0)
if plan.full_ad_fits:
    result = sim.forward(n_steps=10_000)
    report = sim.mesh_intelligence_report(n_steps=10_000, available_memory_gb=24.0)
elif plan.segmented_fits and plan.checkpoint_mode == "checkpoint_segments":
    result = sim.forward(n_steps=10_000, checkpoint_segments=plan.checkpoint_segments)
    report = sim.mesh_intelligence_report(
        n_steps=10_000,
        checkpoint_segments=plan.checkpoint_segments,
        available_memory_gb=24.0,
    )
else:
    raise RuntimeError(plan.recommendation)

From a source checkout, generate the same planning artifact without running FDTD. scripts/ is not installed with the wheel, and --output keeps the artifact out of the documentation tree:

python scripts/memory_reduction_planning_artifact.py \
  --n-steps 10000 \
  --available-memory-gb 8.0 \
  --output /tmp/rfx-memory-plan.json

Use a small loss function or reduced problem for the JAX trace diagnostic:

import jax.numpy as jnp
from rfx import diagnose_ad_saved_residuals

def loss(x):
    return jnp.sum(jnp.sin(x) ** 2)

diagnostic = diagnose_ad_saved_residuals(
    loss,
    jnp.ones((32,)),
    workflow="reduced-inverse-design-loss",
    context={"n_steps": 10_000},
    artifacts={"memory_plan": plan_json},
)
print(diagnostic.parser_health)
print(diagnostic.top_residuals)
print(diagnostic.recommendations)

The report preserves the raw output of jax.ad_checkpoint.print_saved_residuals(...) and adds parsed dtype, shape, byte estimates, original line indices, JAX version, grouped summaries, parser health, optional snapshots, and deterministic recommendations. Use inspect_ad_saved_residuals(...) when only the low-level adapter output is needed.

Use checkify_invariants(...) for bounds or other invariant checks that must survive JAX transforms:

import jax.numpy as jnp
from rfx import check_bounds, checkify_invariants

def loss(eps_r):
    check_bounds(eps_r, lower=1.0, upper=12.0, name="eps_r")
    return jnp.sum(eps_r)

checked_loss = checkify_invariants(loss)
err, value = checked_loss(jnp.ones((16,)))
assert err.get() is None

These checks report runtime invariant failures; they do not assess the RF physics.

Keep the following artifacts with a memory-reduced result:

1. mesh_intelligence_report(...).to_json(),
2. plan_ad_memory(...).to_json() when reverse-mode AD memory is relevant,
3. explain_ad_memory(...).to_json() when the memory breakdown informed the configuration,
4. the exact run command and version or commit,
5. the validation artifact for the reported RF observable,
6. feature-support warnings or validation errors, and
7. gradient evidence when the result depends on jax.grad.