JAX and Flax graphs

TensorViz opens plain JAX functions and Flax Linen/NNX modules through the same Open as Graph, saved setup, Canvas and browser workflow as PyTorch. JAX and Flax are optional model dependencies in the selected capture Python. Missing libraries use the existing reviewed dependency installation flow. The shared capture worker still requires its capture dependencies, including Torch; adding JAX does not replace the selected Torch/CUDA build.

JAX graphs describe a traced program from jax.make_jaxpr, for the saved input shapes, dtypes and static Python values. Shape evidence is inferred, not a measurement of an eager forward. Tracing runs Python and model setup under the existing Workspace Trust, saved-source, selected-interpreter and worker deadline checks. Source-only discovery never imports the model or JAX.

Open a function

Save this as model.py, choose Open as Graph on net, then set x to a float32 tensor with shape 2, 3 in the input form:

import jax
import jax.numpy as jnp

@jax.jit
def net(x):
    return jax.nn.relu(x @ jnp.ones((3, 4)))

The equivalent saved recipe is:

version = 1

[profiles.jax_example]
kind = "model"
framework = "jax"
model = "model:net"
mode = "eval"
device = "cpu"
grad_enabled = false
seed = 0
init = { args = [], kwargs = {} }
inputs = { args = [{ kind = "tensor", shape = [2, 3], dtype = "float32", requires_grad = false, generator = { kind = "full", value = 1 } }], kwargs = {} }

Omitting framework retains existing PyTorch behavior. Plain JAX functions take their arguments from inputs, with an empty init. Array leaves in registered JAX PyTrees are traced; Python scalar arguments remain static. Lists, tuples and dictionaries retain their structure. Flax structs, named tuples, FrozenDicts and other registered containers normalize to dictionaries of their registered child keys. Static PyTree metadata is not serialized. Capture never reconstructs these containers with graph-reference objects, so array invariants in a custom constructor remain valid. This permits an explicitly configured boolean to select Python control flow. Changing any such value requires another trace.

Discovery recognizes top-level functions using imported JAX operations or decorators, and classes deriving from the documented flax.linen.Module and flax.nnx.Module bases, including local inherited/re-exported bases. It does not classify every function in a file merely because that file imports JAX.

Open Flax

from flax import linen as nn

class MLP(nn.Module):
    width: int = 8

    @nn.compact
    def __call__(self, x):
        return nn.Dense(4)(nn.relu(nn.Dense(self.width)(x)))

Open MLP, enter the example input, and retain framework = "jax" in its recipe. Linen constructor annotations populate editable setup fields. An unbound Linen module is initialized with jax.random.key(seed) and the saved inputs. This is freshly initialized state, not trained weights. Its __call__ is then traced.

For saved variables, replace init with a workspace factory returning the exact selected Linen class bound with model.bind(variables). NNX modules can use the same factory mechanism to construct explicit nnx.Rngs, layers and state. A factory for a plain function may return a configured callable. Factories run only inside the approved worker; their files participate in source freshness checks.

Python input functions use the existing inputs.factory contract and can return {"args": (...,), "kwargs": {...}} or a mapping of named inputs. Use these for parameter PyTrees, typed PRNG keys, or loading JAX-specific fixtures. See Python example inputs.

Graph evidence and limits

• Primitive names, argument positions, repeated uses, value shapes/dtypes and nested output structure are retained in the canonical capture. JIT, checkpoint (remat2) and unconditional custom-derivative call bodies expand without executing kernels. These groups describe tracing scopes, not a recovered Flax layer hierarchy.
• cond, scan, while and other unexpanded subprograms retain their primitive boundary and output shapes. Capture coverage is partial with opaque_kernel; neither a branch nor an iteration is reported as executed. Effects are also explicitly incomplete. Convolution/gather dimension tuples, precision and gather-mode enums, and ragged-dot dimension metadata are retained. Other unsupported metadata stays opaque. Scatter reducer subprograms remain explicit opaque boundaries, including in MoE models.
• JAX source frames provide navigation only. PyTorch source-edit recipes, training-step capture, storage aliases, activation statistics and output-byte measurements are unavailable for Jaxpr captures. Large graphs can use Fit below the default minimum zoom; Readable detail returns to legible node cards.
• Setup uses the saved cpu or cuda:N device, eval, and grad_enabled = false. The selected Python must have a compatible JAX CUDA installation for NVIDIA GPUs; an unavailable backend or index fails before model import, without CPU fallback. The device scopes source import, initialization, input factories and tracing. Explicitly placed or sharded arrays from a factory retain their own placement, following JAX's default_device semantics. Graph values remain abstract and do not claim observed GPU execution. Express grad, vmap, or other transformations in the selected callable. JAX has no per-array requires_grad flag. Input descriptions requesting one fail explicitly.
• 64-bit input descriptions require JAX_ENABLE_X64=1 in the capture environment; TensorViz refuses silent dtype narrowing. No Torch checkpoint/fixture loader is used for JAX: use an explicit workspace factory instead.
• Existing operation, call, value, byte, timeout and source-freshness limits apply. Explicit inputs are bounded to 256 arrays and 4,194,304 total elements. Unknown input dimensions are requested rather than guessed. Failed tracing publishes a failure, never a successful empty graph.

NNX modules, bound methods and partial callables clone their state wrappers inside the JAX trace context, allowing nnx.jit to update temporary state without crossing trace levels. Underlying arrays are shared rather than duplicated. Arbitrary Python closures still follow JAX/NNX trace rules.

The capture records backend = "jaxpr", its JAX version and coverage.scope = "traced_program". It cannot carry runtime timings or claim PyTorch execution coverage. Shared graph projection, paginated inspection and portable graph files retain this distinction.