feat[next]: register JaxArrayField as a JAX pytree node - #2732
Conversation
Flatten a JaxArrayField into its backing array (traced child) plus its domain (static aux data), so gt4py fields can be passed to and returned from jax.jit-ed functions outside of field operators.
|
|
||
| result = identity(field) | ||
|
|
||
| assert traced_domains == [domain] # the domain is static metadata, not a traced value |
There was a problem hiding this comment.
What are we actually testing here?
There was a problem hiding this comment.
Not much, honestly — replaced in 2d1173e.
The assertion escaped the trace through a side-effecting list append, so it only held because identity was called exactly once; jax.jit caches by pytree structure and shape/dtype, so a second call or a parametrization would have left traced_domains stale and the failure would have been baffling. Beyond that it mostly duplicated test_jax_jit_field_arguments.
It now flattens and unflattens directly, and asserts the three things the registration is actually promising:
- the single child is the backing array, not the field itself (an unregistered type flattens to one leaf too, so
len(children) == 1alone does not discriminate — this was the accidental part of the old test) - the domain survives
tree_unflatten - a differing domain yields a differing
treedef, which is what makes the domain static and forces a retrace rather than letting a field come back silently mislabelled
Checked that it fails without the registration:
> assert children[0] is field.ndarray
E AssertionError: assert JaxArrayField(...) is Array([[1., 1., 1.], [1., 1., 1.]], dtype=float64)
Replace the jit-passthrough test, whose assertion escaped the trace via a side-effecting list append and therefore only held on the first call, with a flatten/unflatten round-trip that pins the child, the restored domain, and that a differing domain yields a differing tree structure.
There was a problem hiding this comment.
Pull request overview
Registers JaxArrayField as a JAX pytree node so GT4Py fields can be passed into and returned from jax.jit-compiled functions (outside of field-operator contexts) by treating the backing array as the traced leaf and the Domain as static aux data.
Changes:
- Register
JaxArrayFieldwithjax.tree_util.register_pytree_nodeusing(ndarray,)as children anddomainas aux data. - Add unit tests validating
jax.jitargument/return behavior and pytree flatten/unflatten roundtrips.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.
| File | Description |
|---|---|
src/gt4py/next/embedded/nd_array_field.py |
Registers JaxArrayField as a JAX pytree node with custom flatten/unflatten functions. |
tests/next_tests/unit_tests/embedded_tests/test_nd_array_field.py |
Adds JAX-focused tests for jax.jit calls with fields and pytree roundtrip behavior. |
| def _unflatten_jax_field( | ||
| domain: common.Domain, children: tuple[core_defs.NDArrayObject] | ||
| ) -> JaxArrayField: | ||
| return JaxArrayField(domain, children[0]) # type: ignore[abstract] # mypy does not see '__gt_builtin_func__' as implemented by 'FieldBuiltinFuncRegistry' |
There was a problem hiding this comment.
Deliberate — routing through from_array would break more than it fixes.
JAX calls unflatten with things that are not arrays. Measured on this exact registration: object() sentinels and None (vmap, jacfwd), ShapeDtypeStruct (eval_shape), plus four tracer types. from_array rejects all of those, so eval_shape / vmap / jacfwd would start raising.
For the shape-preserving transforms (jit, grad, lax.scan) there is nothing left to validate: the domain is static aux data, resolved at trace time and baked into the very slice constants that produced the array. The mismatch you describe is reachable only under shape-changing transforms, which cannot work with a static domain regardless — validating here would turn "silently wrong" into "raises", not into "works".
| def _unflatten_jax_field( | ||
| domain: common.Domain, children: tuple[core_defs.NDArrayObject] | ||
| ) -> JaxArrayField: | ||
| return JaxArrayField(domain, children[0]) # type: ignore[abstract] # mypy does not see '__gt_builtin_func__' as implemented by 'FieldBuiltinFuncRegistry' |
There was a problem hiding this comment.
Just for my own understanding: shouldn't this distinguish between JaxArrayField and JaxArrayConnectivityField?
Flatten a JaxArrayField into its backing array (traced child) plus its domain (static aux data), so gt4py fields can be passed to and returned from jax.jit-ed functions outside of field operators.