标签：#jax

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