jax_and_decorators

JAX & the decorator system — deep dive

Long-form reference for the grid decorators, the xp (NumPy/JAX) backend pattern, and how autoarray types cross the jax.jit boundary. The per-repo AGENTS.md files keep only a short summary and link here. This is the single canonical source for the detail — PyAutoGalaxy and PyAutoLens point at it rather than re-explaining.

Everything below is grounded in the installed source under autoarray/, autogalaxy/, and autolens/. Where a class or function is named, it exists in the current tree.

---

1. The decorator system

autoarray/structures/decorators/ contains the output-wrapping decorators used on all grid-consuming functions. They ensure the type of the output structure matches the type of the input grid.

Import them as aa.decorators.*. (aa.grid_dec still resolves as a deprecated alias — autoarray/__init__.py defines from .structures import decorators as grid_dec # deprecated alias — but every shipped profile uses aa.decorators.*, so write that form.)

Decorator	Grid2D input →	Grid2DIrregular input →
@aa.decorators.to_array	Array2D	ArrayIrregular
@aa.decorators.to_grid	Grid2D	Grid2DIrregular
@aa.decorators.to_vector_yx	VectorYX2D	VectorYX2DIrregular

How they work

All three share AbstractMaker (decorators/abstract.py). The decorator:

1. Wraps the function in a wrapper(obj, grid, xp=np, *args, **kwargs) signature.
2. Instantiates the relevant *Maker class with the function, object, grid, and xp.
3. AbstractMaker.result checks the grid type and calls the appropriate via_grid_2d / via_grid_2d_irr method to wrap the raw result.

The function body receives the grid as-is and must return a raw array (not an autoarray wrapper). The decorator does the wrapping:

@aa.decorators.to_array
def convergence_2d_from(self, grid, xp=np, **kwargs):
    # grid is Grid2D or Grid2DIrregular — access raw values via grid.array[:, 0]
    y = grid.array[:, 0]
    x = grid.array[:, 1]
    return xp.sqrt(y**2 + x**2)   # return raw array; decorator wraps it