JAX functional backend

ARTI remains PyTorch-first. arti.jax is native JAX code for a deliberately small functional subset—not a second implementation of fitting, attachment, serialization, Recall, or the module stack.

Runtime status and smoke gate

import arti.jax as aj

print(aj.backend_status())          # available | broken | unavailable
report = aj.smoke_report()
print(report["smoke_status"])

The smoke report checks eager execution, JIT, gradients over the whole parameter PyTree, and VMAP semantics. A pass proves only this documented subset.

Functional layer

import jax
import jax.numpy as jnp
import arti.jax as aj

params = aj.init_layer(
    jax.random.PRNGKey(0), input_dim=32, hidden_dim=64, coord_dim=8
)
x = jnp.ones((4, 16, 32))
coord = jnp.ones((4, 16, 8))
mask = jnp.ones((4, 16), dtype=bool)
out = aj.apply_layer(params, x, coord=coord, mask=mask)

params is an array-only PyTree. apply_layer_single() is the explicit per-sample form for composition with jax.vmap.

Shared tensor contracts

The subset includes masked_mean, masked_softmax, mask_coverage, ensure_visibility, attention_mask_to_visibility, and apply_coord_frame_inverse. Tests compare selected mask, visibility, and operator-bank frame-inverse behavior with the PyTorch functional backend.

Boundary

Unsupported here: arti.fit(), attachment, Recall and Stateful Recall, Membrane, Half/Fold/Pulse, checkpointing, multi-device execution, and complete ARTILayer parity. Install-state reporting distinguishes an absent JAX installation from a broken local jaxlib runtime.