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.

Choose a method

Need	Method
Select a checkpointing option for an AD run	Start with ad_memory_preflight(...).
Inspect the components of a static AD estimate	Use explain_ad_memory(...).
Inspect what JAX saves for reverse mode	Use diagnose_ad_saved_residuals(...), or inspect_ad_saved_residuals(...) for the low-level output.
Check the compiler estimate for an exact compiled executable	Use ad_memory_compiled_certificate(...) after compiling that executable.
Reduce cells for a thin feature along one axis	Use a supported non-uniform z mesh, then use checkpoint_every if the AD plan selects it.
Reduce reverse-mode storage on a uniform grid	Use checkpoint_segments if the AD plan selects it.

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.

What the reports measure

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:

evidence_class	Available	Meaning
static_estimate	yes	Formula-based estimate from static metadata such as shape, dtype, NTFF monitors, and active steps.
calibrated_conservative_plan	yes	Budget check built from the static estimates and a safety multiplier.
static_ad_explainability	yes	Breakdown of field, material/CPML, reverse-mode saved-state, and monitor memory in the selected static estimate.
jax_saved_residuals_inspection	yes	Trace-time listing from inspect_ad_saved_residuals(...); it is not an XLA memory report or runtime profile.
rfx_ad_saved_residuals_diagnostic	yes	rfx report from diagnose_ad_saved_residuals(...), including JAX version, parser health, grouped residuals, optional snapshots, and recommendations. It is trace-time information, not profiling.
composite_ad_memory_preflight	yes	Combined budget status, checkpoint selection, static explanation, optional mesh report, optional saved-residual diagnostic, and action hints from ad_memory_preflight(...).
static_action_hint	yes	Deterministic action attached to the preflight report, such as use_checkpoint_every.
bounded_certificate	yes	Compiler memory estimate from ad_memory_compiled_certificate(...) for one exact caller-supplied compiled executable and complete scope metadata.

Run the AD memory preflight

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.

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

Check an exact compiled executable

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.

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.

Configure checkpointing and the mesh

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.

Restricting which cells carry a derivative

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, ...)

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

Inspect saved residuals

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.

Add transformed runtime checks

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.

Archive a run

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.