Skip to content

bug: print_vram_use fails on lmp.Sersic lenses (xp kwarg + JAX tracer) #417

Description

@Jammy2211

Summary

AnalysisImaging.print_vram_use(model=model, batch_size=...) raises on any model whose lens galaxy carries a lmp.Sersic bulge (i.e. a decomposed stellar-mass-via-M/L lens). The failure surfaces on both the use_jax=True and use_jax=False paths, but for different reasons. Both reproduce on canonical main (bc360c08) with no local modifications.

This blocks the modeling.py scripts in:

  • autolens_workspace/scripts/imaging/features/advanced/mass_stellar_dark/modeling.py
  • autolens_workspace/scripts/group/features/advanced/mass_stellar_dark/modeling.py (just shipped in autolens_workspace#164)

The decomposed-mass fit.py / likelihood_function.py scripts in those directories are unaffected because they don't call print_vram_use.

Reproducer

From the autolens_workspace directory on main:

# use_jax=False path
PYAUTO_TEST_MODE=2 PYAUTO_SKIP_CHECKS=1 PYAUTO_WORKSPACE_SMALL_DATASETS=1 \
PYAUTO_DISABLE_JAX=1 PYAUTO_SKIP_WORKSPACE_VERSION_CHECK=1 \
NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib \
python scripts/imaging/features/advanced/mass_stellar_dark/modeling.py

# use_jax=True path (drop PYAUTO_DISABLE_JAX)
PYAUTO_TEST_MODE=2 PYAUTO_SKIP_CHECKS=1 PYAUTO_WORKSPACE_SMALL_DATASETS=1 \
PYAUTO_SKIP_WORKSPACE_VERSION_CHECK=1 \
NUMBA_CACHE_DIR=/tmp/numba_cache MPLCONFIGDIR=/tmp/matplotlib \
python scripts/imaging/features/advanced/mass_stellar_dark/modeling.py

Both fail at the analysis.print_vram_use(model=model, batch_size=search.batch_size) line inside the modeling script.

Branch 1 — use_jax=False: Galaxy.potential_2d_from passes xp to MassProfile.potential_2d_from

File: autogalaxy/galaxy/galaxy.py:379-385

if self.has(cls=MassProfile):
    return sum(
        map(
            lambda p: p.potential_2d_from(grid=grid, xp=xp),   # <-- xp kwarg
            self.cls_list_from(cls=MassProfile),
        )
    )
return xp.zeros((grid.shape[0],))

Error:

TypeError: MassProfile.potential_2d_from() got an unexpected keyword argument 'xp'

MassProfile.potential_2d_from does not accept an xp kwarg. Either:

  • MassProfile.potential_2d_from should accept xp (mirror the kwarg threading used in deflections_yx_2d_from), OR
  • Galaxy.potential_2d_from should stop passing it.

Decomposed-mass lenses surface this because they're the configuration that exercises Galaxy.potential_2d_from from inside print_vram_use (single-mass lenses appear to take a different code path).

Branch 2 — use_jax=True: cse_settings_from calls np.log10 on a traced array

File: autogalaxy/profiles/mass/stellar/sersic.py:56-61

upper_dex = np.min(
    [
        np.log10((23.0 / sersic_constant) ** sersic_index),   # <-- traced sersic_index
        0.85 - np.log10(effective_radius),
    ]
)

Error:

jax.errors.TracerArrayConversionError: The numpy.ndarray conversion method __array__()
was called on traced array with shape float64[]
This BatchTracer with object id ... was created on line:
  autogalaxy/profiles/mass/stellar/sersic.py:58:25 (cse_settings_from)

cse_settings_from is called during the JAX traced compilation of lmp.Sersic's deflection decomposition (via _deflections_2d_via_cse_from). When sersic_index is a JAX traced array, np.log10 cannot operate on it. This needs xp.log10 / xp.min with the analysis's xp module threaded through, OR cse_settings_from needs to receive only concrete numerical inputs (e.g. by treating sersic_index as a static argument at the JAX-jit boundary).

Environment

  • PyAutoGalaxy main at bc360c08
  • PyAutoArray, PyAutoLens, PyAutoFit all on their respective main branches
  • Python 3.12, JAX 0.7 (jax_plugins.xla_cuda12 present but not exercised in CPU repro)

Scope / impact

  • Workspace modeling.py scripts blocked from end-to-end runs for any lmp.Sersic lens (mass_stellar_dark imaging + group, plus any future decomposed-mass tutorial scripts).
  • chaining.py / slam.py in the same directories likely hit the same path through their LIGHT LP → MASS LIGHT DARK stages, but weren't run end-to-end (heavy Nautilus + same library bug).
  • fit.py / likelihood_function.py for decomposed-mass are unaffected — they don't call print_vram_use.

Worth coupling the two fixes since they share the same trigger (lmp.Sersic + print_vram_use) and likely the same module ownership.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions