标签：#jax

This article is about running a hand-written Gemma 4 port in pure JAX on three different accelerators, and about the two places the abstraction leaks. The code is here: github.com/xbill9/gemma4-dev What is this project trying to Do? This project aims to serve one Gemma 4 checkpoint from one JAX port across every accelerator I can rent, and to find out — by measurement, not by reading docs — which parts of "it's just JAX" are true. The port lives in ports/gemma4/ and is driven by a generation loop behind an OpenAI-compatible server. No PyTorch, no vLLM, no torch_xla . The same code runs on Cloud TPU v5e and v6e, and on an NVIDIA T4G attached to an AWS Graviton2 host. "Pure JAX" is the whole experiment. If the port is really portable, the only thing that should change between those rigs is a config file. It mostly is. Two things are not, and they are the interesting part. Gemma 4 E2B is not a stock transformer Any port has to carry four irregularities, and none of them are optional: Two attention geometries. Sliding layers use head_dim=256 , global layers use 512 . Most inference stacks assume one head dimension per model. 8:1 MQA , so the KV budget is nothing like the parameter count would suggest. A KV-share map that collapses 35 layers onto 15 caches . A 512-slot sliding ring , plus per-layer embeddings (PLE) held in a 4.70 GB table that gets quantized to 4 bits on load. That first one is worth dwelling on, because it is what breaks other stacks. On the vLLM path, the heterogeneous head dims force the Triton attention backend: Gemma4 model has heterogeneous head dimensions {'sliding_attention': 256, 'full_attention': 512}. FA4 not available, forcing TRITON_ATTN backend. And on a Turing GPU that backend then asks for shared memory the hardware does not have: triton.runtime.errors.OutOfResources: out of resource: shared memory, Required: 98304, Hardware limit: 65536 JAX never enters that conversation. Attention is ordinary XLA rather than a hand-tiled kernel, so ther

When running large-scale tensor-network contractions with TensorCircuit-NG and the JAX GPU backend, the following runtime configuration is worth testing: XLA_PYTHON_CLIENT_PREALLOCATE = false XLA_FLAGS = --xla_gpu_autotune_level = 0 python your_script.py Its main benefit is not speed, but lower persistent GPU memory usage from XLA GPU autotuning, which makes memory behavior during compilation and on the first visible GPU more predictable. In the TensorCircuit contraction workloads we tested, disabling autotuning also slightly improved steady-state runtime, but the memory savings were the more important result. Key takeaway XLA GPU autotuning evaluates alternative algorithms or workspace configurations for certain GPU kernels and custom calls, then selects an implementation. This can be valuable for convolutions, large GEMMs, and deep-learning workloads with fixed shapes. For large TensorCircuit contractions, however, the contraction path is already determined by OMECO or cotengra, leaving relatively little optimization freedom for autotuning while still potentially incurring substantial persistent memory overhead during compilation and tuning. For TensorCircuit contractions, run this A/B test by default: # Baseline XLA_PYTHON_CLIENT_PREALLOCATE = false python your_script.py # Test configuration XLA_PYTHON_CLIENT_PREALLOCATE = false XLA_FLAGS = --xla_gpu_autotune_level = 0 python your_script.py Both environment variables must be set before the Python process starts and before JAX is imported. This recommendation primarily concerns GPUs; CPU backends do not exhibit the same GPU-kernel autotuning behavior. Representative results All results below use a fixed contraction path so that path-search randomness does not affect the comparison. Workload Autotuning Post-compile memory Peak memory Steady-state runtime 100 qubits × 24 layers, amplitude Default 8.7 GiB 8.7 GiB 0.37 s 100 qubits × 24 layers, amplitude autotune=0 0.5 GiB 4.6 GiB 0.40 s 28 qubits × 12 layers, expecta

PyTrees look deceptively simple. You flatten a nested Python object into leaves, keep a structure descriptor, and later rebuild or map over the same shape. That abstraction is powerful enough to carry optimizer states, model parameters, batched inputs, gradients, and sharding annotations. It is also just ambiguous enough that three major frameworks implement three subtly different languages under the same idea. This note compares JAX jax.tree_util , PyTorch torch.utils._pytree , and TensorFlow tf.nest . I tested the behavior in two environments: an older stack with JAX 0.4.35, PyTorch 2.2.2, TensorFlow 2.20.0, and a newer stack with JAX 0.10.0, PyTorch 2.12.0, TensorFlow 2.21.0. Most flatten/unflatten semantics were stable across these versions. The main version-sensitive result is PyTorch: _pytree.tree_map in 2.2.2 accepts only one pytree, while 2.12.0 supports multiple pytrees and behaves much closer to JAX prefix-style mapping. The short version: JAX treats pytrees as a transformation language, PyTorch is converging toward that model in torch.func , and TensorFlow exposes a broader nested-structure utility through tf.nest . Those differences show up exactly where backend-agnostic libraries usually hurt: None , dictionary order, custom containers, tree_map , autodiff, and vectorization. The Shape Of The APIs The three APIs have the same surface story but not the same contract. from jax import tree_util as jtu from torch.utils import _pytree as tpu import tensorflow as tf leaves , treedef = jtu . tree_flatten ( tree ) tree = jtu . tree_unflatten ( treedef , leaves ) tree = jtu . tree_map ( f , * trees ) leaves , spec = tpu . tree_flatten ( tree ) tree = tpu . tree_unflatten ( leaves , spec ) tree = tpu . tree_map ( f , tree ) # PyTorch 2.2.2 tree = tpu . tree_map ( f , * trees ) # PyTorch 2.12.0 leaves = tf . nest . flatten ( tree ) tree = tf . nest . pack_sequence_as ( structure , leaves ) tree = tf . nest . map_structure ( f , * structures ) Flattening means "whi