From 4550cc8588c4343a45d0cc7d783d5bd6c6eab52b Mon Sep 17 00:00:00 2001 From: Jammy2211 Date: Thu, 3 Sep 2026 18:50:01 -0400 Subject: [PATCH] perf: fold fixed-geometry deflection fields out of the JAX trace (memo phase 2, #604) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit deflections_memo on the JAX backend is now a trace-time constant fold: when the grid is a concrete jax.Array (PyAutoArray evaluates a constant shift-and-rotate at compile time) and every geometry argument is concrete while mass_to_light_ratio is a tracer, the unit-ratio field is evaluated once with numpy/scipy on a numpy twin of the grid, stored in the same memo dict the numpy path uses, and returned as `ratio * xp.asarray(field)`. The Faddeeva subgraph leaves the jaxpr (53,369 -> 13,289 equations on a SLaM-shaped fit, -75%), compile time halves, and the JAX path inherits scipy accuracy for fixed geometry. Concreteness is tested positively (`_is_concrete_array`: ndarray, or a jax.Array that is not a Tracer via sys.modules — never np.asarray in a try); a tracer among the key values or a traced grid falls through to the direct JAX call, so nothing branches on a traced value. `memo_stats()['jax_folds']` is the witness. Numpy behaviour unchanged. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01XhnA4pFN2NycuKc8Ni6s2R --- .../mass/abstract/deflections_memo.py | 343 +++++++++++++++--- .../mass/abstract/test_deflections_memo.py | 108 ++++++ 2 files changed, 406 insertions(+), 45 deletions(-) diff --git a/autogalaxy/profiles/mass/abstract/deflections_memo.py b/autogalaxy/profiles/mass/abstract/deflections_memo.py index 8b7ded71..754e77b4 100644 --- a/autogalaxy/profiles/mass/abstract/deflections_memo.py +++ b/autogalaxy/profiles/mass/abstract/deflections_memo.py @@ -27,9 +27,33 @@ re-masked grid changes the bytes and therefore misses, which is correct by construction. Failure modes are **misses, never stale hits**: anything that cannot be fingerprinted -exactly (a non-scalar constructor argument, a JAX tracer, a grid without a numpy array) -falls through to the ordinary call. The memo engages only when ``xp is np``; the JAX path -is untouched. +exactly (a non-scalar constructor argument, a JAX tracer, a grid that is neither numpy +nor a concrete JAX array) falls through to the ordinary call. + +__The JAX path__ + +On JAX the memo is a **trace-time constant fold** rather than a cross-call cache. Under +the JAX likelihood a fixed Gaussian's geometry reaches the profile as Python floats and +the grid as a *concrete* ``jax.Array`` (``Grid2D.subtracted_and_rotated_from`` evaluates +a constant shift-and-rotate eagerly for exactly this reason), while the free +``mass_to_light_ratio`` is a tracer. The Faddeeva subgraph is therefore a constant, but +``jax.jit`` stages every ``jax.numpy`` call whatever its operands, and this stack +disables XLA's constant folding (``--xla_disable_hlo_passes=constant_folding``, set by +``autonerves/jax_wrapper.py``), so nothing else folds it away either. + +So the memo folds it explicitly: on a miss it evaluates the unit-ratio field **with +numpy and scipy** on a numpy twin of the grid, stores it in the same dict the numpy path +uses (both backends share entries -- the same bytes), and returns +``mass_to_light_ratio * xp.asarray(field)``. The Faddeeva block leaves the jaxpr, which +is replaced by one constant and one multiply, and the JAX path inherits ``scipy.wofz`` +accuracy for fixed geometry rather than the Weideman-32 series. + +The cost is one numpy evaluation per (profile geometry, grid) **per trace**, paid at +compile time. Recompilation is not made more likely by this: the embedded constant +changes only when the geometry or the grid changes, and either of those is a different +fit. Anything traced -- a free geometry parameter, or a free ``grid_offset`` that makes +the grid itself a tracer -- fails the concreteness test and falls through to the ordinary +JAX call, so no branch is ever taken on a traced value. Disable with ``AUTOGALAXY_DEFLECTIONS_MEMO=0``, or in-process with the ``memo_disabled()`` context manager (which a profiling harness needs, since it must be @@ -43,12 +67,19 @@ import hashlib import inspect import os +import sys import weakref import numpy as np from typing import Any, Dict, List, NamedTuple, Optional, Tuple from autoarray.structures.decorators.to_vector_yx import VectorYXMaker +from autoarray.structures.grids.irregular_2d import Grid2DIrregular +from autoarray.structures.grids.uniform_2d import Grid2D + +from autogalaxy.profiles.mass.abstract.mge import ( + _is_static_scalar as _is_concrete_scalar, +) class MemoEntry(NamedTuple): @@ -74,15 +105,16 @@ class MemoEntry(NamedTuple): _DEFAULT_MAX_MB = 256.0 -_stats = {"hits": 0, "misses": 0, "evictions": 0, "bytes": 0} +_stats = {"hits": 0, "misses": 0, "evictions": 0, "bytes": 0, "jax_folds": 0} -# id(grid) -> (weakref to the grid, its fingerprint). One likelihood evaluation calls +# id(grid) -> (weakref to the grid, its fingerprint, its numpy twin or None). One +# likelihood evaluation calls # this module once per mass profile with the *same* grid object, so hashing the grid # bytes per profile would dominate the hit path (~0.2 ms x 30 Gaussians). The weakref # makes a recycled id a miss rather than a wrong answer. It does assume the grid's # coordinate array is not mutated in place after it is first fingerprinted, which no # library path does -- grids are rebuilt, never edited. -_grid_fingerprint_cache: Dict[int, Tuple[Any, str]] = {} +_grid_fingerprint_cache: Dict[int, Tuple[Any, str, Any]] = {} _GRID_FINGERPRINT_CACHE_MAX_ENTRIES = 8 @@ -145,8 +177,13 @@ def memo_max_bytes() -> int: def memo_stats() -> Dict[str, int]: """ - The memo's counters: ``hits``, ``misses``, ``evictions``, ``bytes`` stored and the - number of ``entries``. + The memo's counters: ``hits``, ``misses``, ``evictions``, ``bytes`` stored, the + number of ``entries``, and ``jax_folds``. + + ``jax_folds`` counts the trace-time numpy evaluations performed for a JAX caller -- + one per (profile geometry, grid) per trace. It is the witness that the fold actually + happened: a JAX fit of an MGE with 30 fixed Gaussians reports 30 folds on its + compiling call and none afterwards, and 0 whenever the memo could not key the call. """ return { "hits": _stats["hits"], @@ -154,6 +191,7 @@ def memo_stats() -> Dict[str, int]: "evictions": _stats["evictions"], "bytes": _stats["bytes"], "entries": len(_deflections_memo), + "jax_folds": _stats["jax_folds"], } @@ -163,10 +201,62 @@ def memo_clear() -> None: """ _deflections_memo.clear() _grid_fingerprint_cache.clear() - _stats.update({"hits": 0, "misses": 0, "evictions": 0, "bytes": 0}) + _stats.update({"hits": 0, "misses": 0, "evictions": 0, "bytes": 0, "jax_folds": 0}) + + +def _is_concrete_array(array, xp=np) -> bool: + """ + Whether ``array`` is an array whose contents can be read *now*: a numpy array + always, and on the JAX backend a ``jax.Array`` that is not a tracer. + + The test is positive on both branches, the house style everywhere a traced value + has to be told apart from a concrete one (``mge._is_static_scalar``, + ``autoarray.validate.is_concrete_scalar``). It is never + ``try: np.asarray(...) except``: that reads a tracer's *shape* as success on some + JAX versions and would turn a traced grid into a silently wrong memo entry. + + ``jax`` is looked up in ``sys.modules`` rather than imported. Nothing can be a + ``jax.Array`` in a process that has not imported jax, so the lookup is exact, and a + numpy-only fit never pays an import for a question whose answer is already no. + """ + if isinstance(array, np.ndarray): + return True + + if xp is np: + return False + + if not type(array).__module__.startswith(("jax", "jaxlib")): + return False + + jax = sys.modules.get("jax") + + if jax is None: + return False + + return isinstance(array, jax.Array) and not isinstance(array, jax.core.Tracer) + + +def _concrete_scalar_value(value, xp=np): + """ + The Python value of a concrete 0-d array, or ``None`` if ``value`` is not one. + + A parameter that a JAX fit holds fixed usually arrives as a Python float, but a + pytree-native instance can hand a profile a concrete 0-d ``jax.Array`` instead. That + is as exact a key as the float is, so it is read back with ``float()`` rather than + making the profile unmemoisable. + """ + if not _is_concrete_array(value, xp=xp) or isinstance(value, np.ndarray): + return None + + as_numpy = np.asarray(value) + + if as_numpy.ndim != 0: + return None + + return float(as_numpy) -def _scalar_token(value) -> Optional[str]: +def _scalar_token(value, xp=np) -> Optional[str]: """ An exact, hashable token for a constructor-argument value, or None if the value is not a scalar (or a tuple / list of scalars) and the profile is therefore not @@ -175,22 +265,33 @@ def _scalar_token(value) -> Optional[str]: ``repr`` of a Python float round-trips exactly, so the token distinguishes any two values the deflection calculation would distinguish. A JAX tracer, an ndarray or a nested object all return None, which is what keeps the memo off those profiles. + + A concrete 0-d ``jax.Array`` tokenises as the Python float it holds, so a JAX caller + and a numpy caller of the same fixed geometry land on the *same* key and share one + stored field. """ if value is None or isinstance(value, (bool, str)): return repr(value) if isinstance(value, (int, float)): - return f"{type(value).__name__}:{value!r}" + if isinstance(value, float): + return f"float:{float(value)!r}" + return f"int:{int(value)!r}" if isinstance(value, np.generic): return f"np:{value.dtype.str}:{value.item()!r}" if isinstance(value, (tuple, list)): - tokens = [_scalar_token(item) for item in value] + tokens = [_scalar_token(item, xp=xp) for item in value] if any(token is None for token in tokens): return None return f"{type(value).__name__}({','.join(tokens)})" + concrete = _concrete_scalar_value(value, xp=xp) + + if concrete is not None: + return f"float:{concrete!r}" + return None @@ -228,7 +329,7 @@ def _init_argument_names_from(cls: type) -> Optional[Tuple[str, ...]]: return _init_argument_names[cls] -def _profile_token(profile, skip: Tuple[str, ...] = ()) -> Optional[str]: +def _profile_token(profile, skip: Tuple[str, ...] = (), xp=np) -> Optional[str]: """ An exact token for a profile's identity: its class plus the values of its constructor arguments, read off the instance and excluding the names in ``skip``. @@ -254,7 +355,7 @@ def _profile_token(profile, skip: Tuple[str, ...] = ()) -> Optional[str]: continue if name not in attributes: return None - token = _scalar_token(attributes[name]) + token = _scalar_token(attributes[name], xp=xp) if token is None: return None tokens.append(f"{name}={token}") @@ -262,25 +363,53 @@ def _profile_token(profile, skip: Tuple[str, ...] = ()) -> Optional[str]: return "|".join(tokens) -def _grid_fingerprint(grid) -> Optional[str]: +def _numpy_twin_from(grid, values: np.ndarray): + """ + A numpy grid carrying ``values``, of the same class as ``grid``, or None if that + class is not one the memo knows how to rebuild. + + The twin is what a JAX-backed call is evaluated on: the deflection body reads only + the coordinates, so the twin needs nothing else to be faithful. It is built once per + grid fingerprint and cached beside it. ``Grid2D`` is rebuilt with + ``over_sample_size=1`` deliberately -- the twin is never over-sampled itself (its + caller was already handed the over-sampled coordinates when that is what is being + evaluated), so building a second over-sampled grid inside it would be pure cost. + + An unrecognised grid class returns None, which falls the call through to the direct + JAX evaluation rather than guessing at a constructor. + """ + if isinstance(grid, Grid2D): + return type(grid)(values=values, mask=grid.mask, over_sample_size=1) + + if isinstance(grid, Grid2DIrregular): + return type(grid)(values=values) + + return None + + +def _grid_fingerprint_and_twin(grid, xp=np) -> Tuple[Optional[str], Optional[Any]]: """ - A content fingerprint of ``grid``: the sha256 of its coordinate bytes, its type name, - shape and dtype, and its ``pixel_scales`` / ``origin`` where the type has them. + A content fingerprint of ``grid`` -- the sha256 of its coordinate bytes, its type + name, shape and dtype, and its ``pixel_scales`` / ``origin`` where the type has them + -- together with the numpy twin a JAX-backed grid is evaluated on. - Returns None if the grid carries no numpy coordinate array (a JAX-backed grid, say), - which makes the profile unmemoisable rather than wrongly memoised. + The twin is None for a numpy grid, which is its own twin. A JAX-backed grid is + accepted only when its coordinates are **concrete**: the device-to-host copy that + hashing them costs is paid once per grid object (the weakref cache below), at trace + time. A traced grid, or any other object with no readable coordinates, returns + ``(None, None)`` -- the profile is then not memoisable, rather than wrongly memoised. """ values = getattr(grid, "array", grid) - if not isinstance(values, np.ndarray): - return None + if not _is_concrete_array(values, xp=xp): + return None, None cached = _grid_fingerprint_cache.get(id(grid)) if cached is not None and cached[0]() is grid: - return cached[1] + return cached[1], cached[2] - values = np.ascontiguousarray(values) + values = np.ascontiguousarray(np.asarray(values)) digest = hashlib.sha256() digest.update(type(grid).__name__.encode()) @@ -296,10 +425,20 @@ def _grid_fingerprint(grid) -> Optional[str]: fingerprint = digest.hexdigest() + # A numpy grid is its own twin; holding a strong reference to it here would defeat + # the weakref this cache is keyed by, so only a rebuilt one is stored. + twin = None + + if not isinstance(getattr(grid, "array", grid), np.ndarray): + twin = _numpy_twin_from(grid, values) + + if twin is None: + return None, None + try: reference = weakref.ref(grid) except TypeError: - return fingerprint + return fingerprint, twin if len(_grid_fingerprint_cache) >= _GRID_FINGERPRINT_CACHE_MAX_ENTRIES: dead = [ @@ -308,9 +447,17 @@ def _grid_fingerprint(grid) -> Optional[str]: for key in dead or [next(iter(_grid_fingerprint_cache))]: _grid_fingerprint_cache.pop(key, None) - _grid_fingerprint_cache[id(grid)] = (reference, fingerprint) + _grid_fingerprint_cache[id(grid)] = (reference, fingerprint, twin) + + return fingerprint, twin + - return fingerprint +def _grid_fingerprint(grid, xp=np) -> Optional[str]: + """ + The content fingerprint of ``grid`` alone (see + :func:`_grid_fingerprint_and_twin`). + """ + return _grid_fingerprint_and_twin(grid, xp=xp)[0] def _memo_key(tier: str, profile_token: str, grid_fingerprint: str) -> str: @@ -335,20 +482,85 @@ def _gaussian_class_pair() -> Tuple[type, type]: return _gaussian_classes -def _takes_ratio_split(profile) -> bool: +def _takes_ratio_split(profile, xp=np) -> bool: """ Whether ``profile`` takes the L2 (unit-ratio) split: a ``Gaussian`` mass profile whose deflections are exactly linear in ``mass_to_light_ratio``. ``GaussianGradient`` subclasses ``Gaussian`` but derives its ratio from three parameters, so it is excluded and takes L1. + + On JAX the ratio is accepted when it is a **tracer or a 0-d array** as well as when + it is a concrete scalar: the split exists precisely so a free ratio can multiply a + stored field, and the memo never branches on the ratio's value -- it only multiplies + by it. Anything else (an array of per-pixel ratios, say) is not a scale factor and + is refused. """ gaussian, gaussian_gradient = _gaussian_class_pair() if not isinstance(profile, gaussian) or isinstance(profile, gaussian_gradient): return False - return isinstance(getattr(profile, "mass_to_light_ratio", None), (int, float)) + ratio = getattr(profile, "mass_to_light_ratio", None) + + if _is_concrete_scalar(ratio): + return True + + if xp is np: + return False + + return type(ratio).__module__.startswith(("jax", "jaxlib")) and ( + getattr(ratio, "ndim", None) == 0 + ) + + +def _numpy_profile_from(profile, names: Tuple[str, ...], ratio_split: bool, xp=np): + """ + A shallow copy of ``profile`` whose every constructor argument is a plain Python + value, ready to be evaluated with numpy, or None if any of them is not. + + This is the object the trace-time fold evaluates. Its whole purpose is that it + carries **no tracer anywhere**: a 0-d concrete ``jax.Array`` parameter is read back + as a float, and the final loop re-checks every argument rather than trusting that + the token pass covered it -- a tracer smuggled onto the copy would be evaluated by + numpy into an array of the wrong thing, silently. + + For an L2 profile the copy's ``mass_to_light_ratio`` is 1.0, so what is computed and + stored is the unit-ratio field the caller then scales by the real (traced) ratio. + """ + profile_np = copy.copy(profile) + + if ratio_split: + profile_np.mass_to_light_ratio = 1.0 + + attributes = vars(profile_np) + + for name in names: + if name not in attributes: + return None + + value = attributes[name] + concrete = _concrete_scalar_value(value, xp=xp) + + if concrete is not None: + setattr(profile_np, name, concrete) + elif isinstance(value, (tuple, list)): + items = [_concrete_scalar_value(item, xp=xp) for item in value] + if any(item is not None for item in items): + setattr( + profile_np, + name, + type(value)( + original if item is None else item + for original, item in zip(value, items) + ), + ) + + for name in names: + if _scalar_token(vars(profile_np)[name], xp=np) is None: + return None + + return profile_np def _unwrap(result) -> Tuple[Optional[np.ndarray], bool]: @@ -414,6 +626,13 @@ def deflections_yx_2d_from(profile, grid, xp=np): This is the single intercept the summation call sites (``Galaxy`` and ``Basis``) go through, so no mass profile needs an override of its own. + On numpy this is a cross-evaluation cache; on JAX it is a trace-time constant fold, + evaluated with numpy on a twin of the (concrete) grid and returned as a constant the + traced ratio multiplies. Both write the same bytes into the same dict. A tracer + among the key values, or a traced grid, takes neither path: the call falls straight + through to ``profile.deflections_yx_2d_from``, so nothing here ever branches on a + traced value. + Parameters ---------- profile @@ -421,25 +640,25 @@ def deflections_yx_2d_from(profile, grid, xp=np): grid The 2D (y,x) coordinates the deflection angles are computed on. xp - The array backend. The memo engages only for numpy; JAX falls straight through. + The array backend, numpy or ``jax.numpy``. """ def direct(): return profile.deflections_yx_2d_from(grid=grid, xp=xp) - if xp is not np or not memo_enabled(): + if not memo_enabled(): return direct() - ratio_split = _takes_ratio_split(profile) + ratio_split = _takes_ratio_split(profile, xp=xp) - profile_token = _profile_token( - profile, skip=("mass_to_light_ratio",) if ratio_split else () - ) + skip = ("mass_to_light_ratio",) if ratio_split else () + + profile_token = _profile_token(profile, skip=skip, xp=xp) if profile_token is None: return direct() - grid_fingerprint = _grid_fingerprint(grid) + grid_fingerprint, twin = _grid_fingerprint_and_twin(grid, xp=xp) if grid_fingerprint is None: return direct() @@ -452,29 +671,63 @@ def direct(): if entry is not None: _stats["hits"] += 1 + field = entry.field if xp is np else xp.asarray(entry.field) if ratio_split: return _wrapped_result( - profile.mass_to_light_ratio * entry.field, grid, xp, entry.wrapped + profile.mass_to_light_ratio * field, grid, xp, entry.wrapped ) - return _wrapped_result(np.array(entry.field), grid, xp, entry.wrapped) + return _wrapped_result( + np.array(entry.field) if xp is np else field, grid, xp, entry.wrapped + ) _stats["misses"] += 1 - if ratio_split: - unit_profile = copy.copy(profile) - unit_profile.mass_to_light_ratio = 1.0 - result = unit_profile.deflections_yx_2d_from(grid=grid, xp=xp) - else: - result = direct() + if xp is np: + if ratio_split: + unit_profile = copy.copy(profile) + unit_profile.mass_to_light_ratio = 1.0 + result = unit_profile.deflections_yx_2d_from(grid=grid, xp=xp) + else: + result = direct() + + field, wrapped = _unwrap(result) + + if field is None: + return result + + _store(key, field, wrapped) + + if ratio_split: + return _wrapped_result( + profile.mass_to_light_ratio * field, grid, xp, wrapped + ) + + return result + + # The JAX fold: evaluate the field once, now, with numpy and scipy on the numpy twin + # of the grid, and hand the trace a constant. + names = _init_argument_names_from(type(profile)) + + profile_np = _numpy_profile_from( + profile, names=names, ratio_split=ratio_split, xp=xp + ) + + if profile_np is None: + return direct() + + result = profile_np.deflections_yx_2d_from(grid=twin, xp=np) field, wrapped = _unwrap(result) if field is None: - return result + return direct() _store(key, field, wrapped) + _stats["jax_folds"] += 1 + + field = xp.asarray(field) if ratio_split: return _wrapped_result(profile.mass_to_light_ratio * field, grid, xp, wrapped) - return result + return _wrapped_result(field, grid, xp, wrapped) diff --git a/test_autogalaxy/profiles/mass/abstract/test_deflections_memo.py b/test_autogalaxy/profiles/mass/abstract/test_deflections_memo.py index 7c642433..0b067087 100644 --- a/test_autogalaxy/profiles/mass/abstract/test_deflections_memo.py +++ b/test_autogalaxy/profiles/mass/abstract/test_deflections_memo.py @@ -1,3 +1,4 @@ +import sys import numpy as np import pytest @@ -193,6 +194,7 @@ def test__kill_switch_disables_the_memo(monkeypatch): "evictions": 0, "bytes": 0, "entries": 0, + "jax_folds": 0, } assert values(result) == pytest.approx( values(profile.deflections_yx_2d_from(grid=grid)), rel=1e-12, abs=0.0 @@ -286,6 +288,7 @@ def test__profile_holding_a_non_scalar_attribute_falls_through(): "evictions": 0, "bytes": 0, "entries": 0, + "jax_folds": 0, } assert result is not None @@ -346,3 +349,108 @@ def basis(mass_to_light_ratio): assert deflections_memo.memo_stats()["hits"] == 3 assert deflections_memo.memo_stats()["misses"] == 3 assert values(on) == pytest.approx(values(off), rel=1e-12, abs=0.0) + + +# --------------------------------------------------------------------------- +# Phase 2 -- the JAX trace-time fold, tested numpy-side only. +# +# `test_autogalaxy` never imports jax, so these exercise the parts of the fold that +# are backend-agnostic: the concreteness test, the numpy twin, and the guarantee that +# a numpy caller never enters the fold at all. +# --------------------------------------------------------------------------- + + +class _FakeTracer: + """ + A stand-in for a JAX tracer: an object whose *type* is defined in `jax._src.core`, + which is how the memo tells a traced value apart without importing jax. + """ + + ndim = 2 + + +_FakeTracer.__module__ = "jax._src.core" + + +def test__is_concrete_array__ndarray_yes__list_and_tracer_no(): + assert deflections_memo._is_concrete_array(np.zeros((3, 2)), xp=np) is True + + # Not an array at all. + assert deflections_memo._is_concrete_array([[0.0, 0.0]], xp=np) is False + assert deflections_memo._is_concrete_array(3.0, xp=np) is False + + tracer = _FakeTracer() + + # A jax-typed object is rejected on the numpy backend by the `xp is np` gate, and on + # a non-numpy backend by the `isinstance(..., jax.Array)` test (or, in a process + # that never imported jax, by the `sys.modules` lookup failing outright). Neither + # route imports jax itself -- other tests in this suite build `use_jax` analyses, so + # whether jax is already loaded is not this test's to assert, only that this call + # does not change it. + jax_was_imported = "jax" in sys.modules + + assert deflections_memo._is_concrete_array(tracer, xp=np) is False + assert deflections_memo._is_concrete_array(tracer, xp=object()) is False + assert ("jax" in sys.modules) is jax_was_imported + + +def test__numpy_twin_reproduces_the_coordinates_and_the_mask(): + grid = grid_2d() + + twin = deflections_memo._numpy_twin_from(grid, np.array(grid.array)) + + assert type(twin) is type(grid) + assert twin.array == pytest.approx(np.array(grid.array), rel=0.0, abs=0.0) + assert (np.asarray(twin.mask) == np.asarray(grid.mask)).all() + assert twin.pixel_scales == grid.pixel_scales + + irregular = ag.Grid2DIrregular(values=np.array(grid.array)) + + twin = deflections_memo._numpy_twin_from(irregular, np.array(irregular.array)) + + assert type(twin) is type(irregular) + assert twin.array == pytest.approx(np.array(irregular.array), rel=0.0, abs=0.0) + + # A grid class the memo cannot rebuild is not guessed at. + assert deflections_memo._numpy_twin_from(object(), np.zeros((2, 2))) is None + + +def test__numpy_backend_never_folds__jax_folds_counter_stays_zero(): + grid = grid_2d() + profile = gaussian() + + for _ in range(3): + deflections_memo.deflections_yx_2d_from(profile, grid, np) + + stats = deflections_memo.memo_stats() + + assert stats["misses"] == 1 + assert stats["hits"] == 2 + assert stats["jax_folds"] == 0 + + # An L1 profile too, and the fingerprint helper's twin is None for numpy -- a numpy + # grid is its own twin, and holding it here would defeat the weakref cache. + isothermal = ag.mp.Isothermal(centre=(0.0, 0.0), einstein_radius=1.0) + deflections_memo.deflections_yx_2d_from(isothermal, grid, np) + + assert deflections_memo.memo_stats()["jax_folds"] == 0 + + fingerprint, twin = deflections_memo._grid_fingerprint_and_twin(grid, xp=np) + + assert fingerprint is not None + assert twin is None + + +def test__concrete_scalar_value__reads_back_nothing_on_the_numpy_backend(): + assert deflections_memo._concrete_scalar_value(1.5, xp=np) is None + assert deflections_memo._concrete_scalar_value(np.array(1.5), xp=np) is None + assert deflections_memo._concrete_scalar_value(_FakeTracer(), xp=np) is None + + # A float and a 0-d value of the same number share one token, so a numpy and a JAX + # caller of the same fixed geometry land on the same memo entry. `np.float64` + # subclasses `float`, so it shares that token too; a `np.float32` does not and keeps + # its own dtype-tagged token, because it is a different number. + assert deflections_memo._scalar_token(1.5) == "float:1.5" + assert deflections_memo._scalar_token(np.float64(1.5)) == "float:1.5" + assert deflections_memo._scalar_token(np.float32(1.5)) == "np: