Backends and integrations

PyTorch applications

PyTorch is a core package dependency. Native models, attached programs and documented pretrained workflows can reach the same ARTITrainingSession. Select a CPU or CUDA Torch build for the actual application rather than inferring hardware readiness from the ARTI package version.

Native JAX programs

arti.jax binds the shared Formula and bounded/static ProgramGraph declarations to native arrays. The static program, parameter PyTree and forward-state PyTree remain separate. Compiled invocation runs JAX operations without executing Torch or a host callback.

python -m pip install "arti-fit[jax,flax]==3.3.0"

• Use jit, grad and vmap within the binding's documented contracts.
• Fixed-horizon scan carries explicit returned state.
• Physical shape and graph structure specialize a binding; custom Torch callables or VJPs require explicit native implementations.

Optional Flax integration

Native Flax integration follows init/apply and separates parameters from mutable collections. Declared mutable state returns in the next state rather than becoming an implicit optimizer parameter. Native attachment and parameter masks let an application retain control over host and adapter ownership.

Keep each adapter's scope

Accelerate, Transformers/TRL, Ray and reporting adapters connect to their documented native workflows. Their upstream loss, accumulation, masks, precision and execution semantics remain relevant. The existence of an adapter does not qualify every model or training topology.

Read the native lowering limits

The documented JAX scope does not include predecessor-writeback Formula effects, automatic lowering of arbitrary federated Query programs, open-horizon reverse mode or automatic independent optimizer-region derivation. Local training recovery does not imply native JAX or distributed/sharded checkpoint recovery.