# Reduce Memory Use

Channel: release; package: 1.8.0; source: 7d1e97331656dd74ad0f5471653ac0c5effdff70

Page: https://remilab.ai/rfx/versions/v1.8.0/guide/memory-reduction/
Source: https://github.com/bk-squared/rfx/blob/7d1e97331656dd74ad0f5471653ac0c5effdff70/docs/public/guide/memory-reduction.mdx

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.

Setup used by the snippets below:

```python
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")
```

```python
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`.

```python
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.

```python
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.

```python
# `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](https://remilab.ai/rfx/versions/v1.8.0/guide/nonuniform-mesh/) to check supported combinations and mesh
restrictions.

```python
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:

```python
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:

```python
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:

```python
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:

```bash
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:

```python
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:

```python
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`.