<?xml version="1.0" encoding="UTF-8"?>
<rss version="2.0" xmlns:atom="http://www.w3.org/2005/Atom" xmlns:dc="http://purl.org/dc/elements/1.1/">
  <channel>
    <title>DEV Community: Shixin Zhang</title>
    <description>The latest articles on DEV Community by Shixin Zhang (@refractionray).</description>
    <link>https://dev.to/refractionray</link>
    <image>
      <url>https://media2.dev.to/dynamic/image/width=90,height=90,fit=cover,gravity=auto,format=auto/https:%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Fuser%2Fprofile_image%2F3763205%2F5c5af020-0b22-4443-aa68-28b3150f48e4.png</url>
      <title>DEV Community: Shixin Zhang</title>
      <link>https://dev.to/refractionray</link>
    </image>
    <atom:link rel="self" type="application/rss+xml" href="https://dev.to/feed/refractionray"/>
    <language>en</language>
    <item>
      <title>From Parameter Tuning to Cross-Paradigm Exploration: Quantum Control Enters the Era of “Autopilot”</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Tue, 21 Jul 2026 06:37:46 +0000</pubDate>
      <link>https://dev.to/refractionray/from-parameter-tuning-to-cross-paradigm-exploration-quantum-control-enters-the-era-of-autopilot-5bep</link>
      <guid>https://dev.to/refractionray/from-parameter-tuning-to-cross-paradigm-exploration-quantum-control-enters-the-era-of-autopilot-5bep</guid>
      <description>&lt;p&gt;If you have ever tried to drive a high-performance race car on ice, you may understand the frustration researchers face when controlling quantum systems today.&lt;/p&gt;

&lt;p&gt;In the grand vision of quantum computing, &lt;strong&gt;Quantum Optimal Control (QOC)&lt;/strong&gt; serves as the steering wheel that guides quantum systems toward desired states. Whether in adiabatic quantum computation, quantum annealing, or quantum simulation, the fundamental challenge remains the same: designing time-dependent control protocols that drive a quantum system from a simple initial state to a complex target state with high fidelity.&lt;/p&gt;

&lt;p&gt;However, real-world quantum control faces two fundamental obstacles.&lt;/p&gt;

&lt;p&gt;On the hardware side, quantum systems are extremely fragile. Short coherence times, limited control channels, hardware-specific constraints, and strict pulse boundaries severely restrict the available control space.&lt;/p&gt;

&lt;p&gt;On the algorithmic side, many-body Hamiltonians associated with practical problems often exhibit complicated spectral structures, including small energy gaps and rugged optimization landscapes. Finding an efficient evolution path within a limited time window remains highly challenging.&lt;/p&gt;

&lt;p&gt;For decades, designing quantum control protocols has largely remained a &lt;strong&gt;human-driven, handcrafted process&lt;/strong&gt;. Experts repeatedly design, simulate, and tune protocols for specific physical systems and hardware platforms through extensive trial and error.&lt;/p&gt;

&lt;p&gt;A recent work introduces a fundamentally different approach: &lt;strong&gt;QOC-Workbench&lt;/strong&gt;, an LLM-driven, fully auditable framework for cross-paradigm quantum control design. Rather than acting as another black-box optimizer, it functions more like an &lt;strong&gt;autopilot system for quantum control&lt;/strong&gt;—transforming how control protocols are discovered, validated, and improved.&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;Reference:&lt;br&gt;
 &lt;em&gt;LLM-Driven Cross-Paradigm Design for Quantum Optimal Control&lt;/em&gt;&lt;br&gt;
 Yu-Qin Chen and Shi-Xin Zhang, arXiv:2607.17498&lt;/p&gt;
&lt;/blockquote&gt;




&lt;p&gt;&lt;a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fwr8ir74r1nc54fvmnf1f.png" class="article-body-image-wrapper"&gt;&lt;img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fwr8ir74r1nc54fvmnf1f.png" alt="QOC-Workbench Architecture" width="800" height="365"&gt;&lt;/a&gt;&lt;/p&gt;

&lt;h1&gt;
  
  
  Beyond Closed Optimization Spaces: How QOC-Workbench Works
&lt;/h1&gt;

&lt;p&gt;Traditional quantum optimal control methods usually operate inside a predefined design space.&lt;/p&gt;

&lt;p&gt;Researchers first choose a control ansatz—a mathematical form for pulse schedules, interpolation functions, or auxiliary Hamiltonians. Classical optimization algorithms then search for optimal parameters within that fixed structure.&lt;/p&gt;

&lt;p&gt;This approach is powerful, but fundamentally limited.&lt;/p&gt;

&lt;p&gt;If the initial design space is incomplete, optimization can only find the best solution &lt;strong&gt;within existing assumptions&lt;/strong&gt;. It cannot invent new functional forms, discover alternative control mechanisms, or challenge the original modeling choices.&lt;/p&gt;

&lt;p&gt;QOC-Workbench changes this paradigm by integrating:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;large language models with scientific reasoning capabilities,&lt;/li&gt;
&lt;li&gt;structured knowledge extracted from quantum control literature,&lt;/li&gt;
&lt;li&gt;historical simulation results,&lt;/li&gt;
&lt;li&gt;hardware constraints,&lt;/li&gt;
&lt;li&gt;and high-performance quantum simulation infrastructure.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;The workflow forms a closed-loop scientific discovery system:&lt;/p&gt;

&lt;h3&gt;
  
  
  1. Human experts define the physical boundary
&lt;/h3&gt;

&lt;p&gt;Researchers specify:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;target Hamiltonians,&lt;/li&gt;
&lt;li&gt;hardware limitations,&lt;/li&gt;
&lt;li&gt;physical constraints,&lt;/li&gt;
&lt;li&gt;evaluation objectives.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;Humans provide the scientific context and ensure physical validity.&lt;/p&gt;

&lt;h3&gt;
  
  
  2. LLM performs cross-paradigm exploration
&lt;/h3&gt;

&lt;p&gt;Instead of only optimizing parameters, the LLM can:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;propose new control schedule families,&lt;/li&gt;
&lt;li&gt;modify the structure of auxiliary Hamiltonians,&lt;/li&gt;
&lt;li&gt;combine ideas from different control paradigms,&lt;/li&gt;
&lt;li&gt;generate executable simulation code.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;The search space itself becomes dynamic.&lt;/p&gt;

&lt;h3&gt;
  
  
  3. Physics solvers provide rigorous validation
&lt;/h3&gt;

&lt;p&gt;Candidate protocols are evaluated through differentiable quantum many-body simulations powered by high-performance quantum software infrastructure such as &lt;strong&gt;TensorCircuit-NG&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;The system does not rely on language-model judgment alone—the generated ideas must survive quantitative physical evaluation.&lt;/p&gt;

&lt;h3&gt;
  
  
  4. Memory engine turns experiments into reusable knowledge
&lt;/h3&gt;

&lt;p&gt;Every successful or failed experiment is automatically recorded as structured knowledge.&lt;/p&gt;

&lt;p&gt;Over time, the system accumulates reusable design principles, allowing previous discoveries to influence future exploration.&lt;/p&gt;

&lt;p&gt;Through this process, AI evolves from a parameter fitting tool into a continuously improving scientific assistant.&lt;/p&gt;




&lt;h1&gt;
  
  
  Three Levels of Evolution: From Pulse Shaping to Neural Control Generators
&lt;/h1&gt;

&lt;p&gt;To demonstrate the capability of QOC-Workbench, the authors tested it across three increasingly challenging physical scenarios.&lt;/p&gt;

&lt;p&gt;These examples illustrate a gradual transition:&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;from optimizing existing protocols → modifying physical pathways → discovering new computational paradigms.&lt;/strong&gt;&lt;/p&gt;




&lt;h2&gt;
  
  
  Case 1: Designing Hardware-Compatible Control Pulses
&lt;/h2&gt;

&lt;h3&gt;
  
  
  Rydberg Atom Arrays
&lt;/h3&gt;

&lt;p&gt;The first challenge considers solving the Maximum Independent Set problem using Rydberg atom arrays.&lt;/p&gt;

&lt;p&gt;Real quantum hardware imposes strict constraints on available control signals. Traditional approaches often rely on analytical counterdiabatic protocols derived from simplified models.&lt;/p&gt;

&lt;p&gt;However, these idealized solutions may not fully capture the complexity of interacting many-body systems.&lt;/p&gt;

&lt;p&gt;QOC-Workbench analyzed the limitations of existing approaches and explored a broader control space.&lt;/p&gt;

&lt;p&gt;Instead of simply tuning parameters of known pulses, it discovered a new pulse structure:&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;the “smooth beta-bump” envelope.&lt;/strong&gt;&lt;/p&gt;

&lt;p&gt;The generated protocol:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;satisfies realistic hardware constraints,&lt;/li&gt;
&lt;li&gt;preserves smooth control behavior,&lt;/li&gt;
&lt;li&gt;achieves higher ground-state fidelity than classical analytical baselines in many-body simulations.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;This demonstrates that LLM-driven exploration can redesign control waveforms rather than merely optimize them.&lt;/p&gt;

&lt;p&gt;&lt;a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fqyta87ehcrrne5dxnk4g.png" class="article-body-image-wrapper"&gt;&lt;img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fqyta87ehcrrne5dxnk4g.png" alt="Agent exploration history" width="799" height="541"&gt;&lt;/a&gt;&lt;/p&gt;




&lt;h2&gt;
  
  
  Case 2: Redesigning Many-Body Evolution Paths
&lt;/h2&gt;

&lt;h3&gt;
  
  
  XXZ Spin Chains
&lt;/h3&gt;

&lt;p&gt;The second example moves beyond pulse engineering.&lt;/p&gt;

&lt;p&gt;For XXZ spin chains with complex spectral structures, QOC-Workbench explored the structure of the entire evolution trajectory.&lt;/p&gt;

&lt;p&gt;The system introduced:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;a target catalyst Hamiltonian that asymptotically vanishes at the endpoints,&lt;/li&gt;
&lt;li&gt;jointly optimized approximate counterdiabatic corrections,&lt;/li&gt;
&lt;li&gt;nonlinear time schedules combining polynomial and sinusoidal deformations.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;The key insight was that optimal control is not only about adjusting individual parameters.&lt;/p&gt;

&lt;p&gt;The &lt;strong&gt;geometry of the evolution path itself&lt;/strong&gt; can be redesigned.&lt;/p&gt;

&lt;p&gt;By jointly optimizing the Hamiltonian pathway and correction terms, the system discovered improved protocols that would be difficult to obtain through conventional parameter optimization alone.&lt;/p&gt;




&lt;h2&gt;
  
  
  Case 3: Breaking the Scaling Barrier with Neural Generators
&lt;/h2&gt;

&lt;h3&gt;
  
  
  2D Random-Field Ising Models
&lt;/h3&gt;

&lt;p&gt;The third case reveals the most significant conceptual shift.&lt;/p&gt;

&lt;p&gt;For disordered many-body systems, optimizing a control protocol separately for every instance quickly becomes computationally expensive.&lt;/p&gt;

&lt;p&gt;QOC-Workbench identified this bottleneck and changed the problem formulation.&lt;/p&gt;

&lt;p&gt;Instead of searching for an optimal protocol instance by instance, it designed and trained a &lt;strong&gt;graph neural network (GNN) generator&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;The generator was trained only on small-scale graph instances but successfully generalized to larger unseen systems.&lt;/p&gt;

&lt;p&gt;It could:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;accurately predict control coefficients,&lt;/li&gt;
&lt;li&gt;generate reasonable evolution paths,&lt;/li&gt;
&lt;li&gt;bypass expensive per-instance variational optimization.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;This represents a transition from:&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;“optimize every problem separately”&lt;/p&gt;
&lt;/blockquote&gt;

&lt;p&gt;to:&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;“learn the underlying structure of the solution space.”&lt;/p&gt;
&lt;/blockquote&gt;




&lt;h1&gt;
  
  
  Conclusion: Let Physicists Return to Physics
&lt;/h1&gt;

&lt;p&gt;QOC-Workbench is not designed to replace human scientific intuition.&lt;/p&gt;

&lt;p&gt;Instead, it aims to amplify it.&lt;/p&gt;

&lt;p&gt;In this emerging human-AI collaboration paradigm, researchers no longer need to spend most of their time on:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;manual parameter tuning,&lt;/li&gt;
&lt;li&gt;repetitive protocol benchmarking,&lt;/li&gt;
&lt;li&gt;low-level implementation details.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;Instead, they can focus on:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;understanding fundamental physical mechanisms,&lt;/li&gt;
&lt;li&gt;defining meaningful physical constraints,&lt;/li&gt;
&lt;li&gt;interpreting machine-discovered protocols,&lt;/li&gt;
&lt;li&gt;extracting new scientific principles.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;Human insights then become new knowledge injected back into the system, creating a continuous feedback loop between human reasoning and machine exploration.&lt;/p&gt;

&lt;p&gt;From manually crafting isolated interpolation curves to building a continuously evolving, auditable, and transferable knowledge system, quantum control is moving beyond fixed optimization frameworks.&lt;/p&gt;

&lt;p&gt;The future of quantum control may not be about finding better parameters inside predefined spaces.&lt;/p&gt;

&lt;p&gt;It may be about building intelligent systems capable of discovering entirely new control paradigms.&lt;/p&gt;

&lt;p&gt;Quantum control is entering the era of &lt;strong&gt;autopilot&lt;/strong&gt;.&lt;/p&gt;

</description>
      <category>ai</category>
      <category>agents</category>
      <category>quantum</category>
    </item>
    <item>
      <title>Training a 1,000-Qubit, 40,000-Parameter Quantum Algorithm with Full Gradients Using TensorCircuit-NG</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Sat, 18 Jul 2026 02:47:36 +0000</pubDate>
      <link>https://dev.to/refractionray/training-a-1000-qubit-40000-parameter-quantum-algorithm-with-full-gradients-using-4pho</link>
      <guid>https://dev.to/refractionray/training-a-1000-qubit-40000-parameter-quantum-algorithm-with-full-gradients-using-4pho</guid>
      <description>&lt;h2&gt;
  
  
  Introduction
&lt;/h2&gt;

&lt;p&gt;If you wanted to simulate 1,000 quantum qubits and compute exact full gradients for 40,000 parameters on a classical computer, how much compute power would you need?&lt;/p&gt;

&lt;p&gt;For a 1,000-qubit, 10-layer circuit, calculating a single state amplitude might have a manageable computational overhead. However, when you attempt to solve the system's Variational Quantum Eigensolver (VQE) and compute all of its parameter gradients, the difficulty skyrockets. At first glance, this sounds like a job that strictly requires a supercomputing cluster.&lt;/p&gt;

&lt;p&gt;The underlying reasons for this explosion in complexity are twofold:&lt;/p&gt;

&lt;ol&gt;
&lt;li&gt;
&lt;strong&gt;The Introduction of the Hamiltonian&lt;/strong&gt;: Computing the VQE expectation value requires evaluating $E(\boldsymbol\theta)=\langle\psi(\boldsymbol\theta)\vert H\vert\psi(\boldsymbol\theta)\rangle$. This transforms the tensor network from a single-sided state contraction into a "bra-MPO-ket" three-layer sandwich structure, immediately more than doubling the graph size.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;The Memory Wall in Backpropagation&lt;/strong&gt;: Deriving gradients for tens of thousands of parameters requires backpropagation. To compute these gradients in a contraction graph, the compiler has to keep massive amounts of intermediate results from the forward pass in VRAM. This is vastly more difficult—and puts far more pressure on memory—than a simple forward contraction or amplitude calculation.&lt;/li&gt;
&lt;/ol&gt;

&lt;p&gt;In this technical blog, we'll demonstrate how we cut the per-step execution time of a 1,000-qubit VQE down to just &lt;strong&gt;1 second on a single NVIDIA RTX 6000D GPU&lt;/strong&gt;. We achieve this by exploring three different physics- and engineering-driven contraction strategies, all without sacrificing exact gradient precision.&lt;/p&gt;

&lt;h2&gt;
  
  
  Performance Comparison: Three Physical Perspectives
&lt;/h2&gt;

&lt;p&gt;We use a 1D Transverse Field Ising Model (TFIM) with open boundary conditions as our benchmark. Each layer of the circuit applies diagonal entangling gates sequentially to adjacent pairs $(0,1),(1,2),\ldots,(n-2,n-1)$ along the qubit chain:&lt;/p&gt;

&lt;p&gt;$$R_{ZZ}(\theta)=\exp(-i\theta Z\otimes Z/2)$$&lt;/p&gt;

&lt;p&gt;Followed by single-qubit rotations. The TFIM Hamiltonian is:&lt;/p&gt;

&lt;p&gt;$$H_{\rm TFIM}=\sum_{i=0}^{n-2}X_iX_{i+1}+\sum_{i=0}^{n-1}Z_i$$&lt;/p&gt;

&lt;p&gt;We can also represent this Hamiltonian as an exact and compact Matrix Product Operator (MPO, with a Bond Dimension of 3).&lt;/p&gt;

&lt;p&gt;The table below shows the benchmark results for three different contraction modes (using an NVIDIA RTX 6000D at &lt;code&gt;complex64&lt;/code&gt; precision):&lt;/p&gt;

&lt;div class="table-wrapper-paragraph"&gt;&lt;table&gt;
&lt;thead&gt;
&lt;tr&gt;
&lt;th&gt;&lt;strong&gt;Perspective (Algorithm Mode)&lt;/strong&gt;&lt;/th&gt;
&lt;th&gt;&lt;strong&gt;Initial Compilation&lt;/strong&gt;&lt;/th&gt;
&lt;th&gt;&lt;strong&gt;First Full Gradient&lt;/strong&gt;&lt;/th&gt;
&lt;th&gt;&lt;strong&gt;Steady-State Full Gradient&lt;/strong&gt;&lt;/th&gt;
&lt;th&gt;&lt;strong&gt;Peak VRAM&lt;/strong&gt;&lt;/th&gt;
&lt;th&gt;&lt;strong&gt;Applicability &amp;amp; Value&lt;/strong&gt;&lt;/th&gt;
&lt;/tr&gt;
&lt;/thead&gt;
&lt;tbody&gt;
&lt;tr&gt;
&lt;td&gt;&lt;strong&gt;Global Contraction Tree&lt;/strong&gt;&lt;/td&gt;
&lt;td&gt;5 h 20 min&lt;/td&gt;
&lt;td&gt;5.88 s&lt;/td&gt;
&lt;td&gt;1.09 s&lt;/td&gt;
&lt;td&gt;51.97 GB&lt;/td&gt;
&lt;td&gt;Directly contracts the full bra-MPO-ket graph; useful for full-graph compute and cross-benchmarking.&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;&lt;strong&gt;Local Causal Sliding Window&lt;/strong&gt;&lt;/td&gt;
&lt;td&gt;10.79 s&lt;/td&gt;
&lt;td&gt;148.68 s&lt;/td&gt;
&lt;td&gt;148.36 s&lt;/td&gt;
&lt;td&gt;8.86 GB&lt;/td&gt;
&lt;td&gt;Relies on strict causal cones of fixed width for local terms; terms can be dispatched independently.&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;&lt;strong&gt;Local Window (8-GPU Parallel)&lt;/strong&gt;&lt;/td&gt;
&lt;td&gt;10.50–11.28 s / GPU&lt;/td&gt;
&lt;td&gt;19.72–20.50 s&lt;/td&gt;
&lt;td&gt;18.71 s&lt;/td&gt;
&lt;td&gt;8.86 GB / GPU&lt;/td&gt;
&lt;td&gt;Dispatches 125 local terms per GPU; yielded a measured 7.93× speedup.&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;&lt;strong&gt;Spatial Transfer Matrix (B3)&lt;/strong&gt;&lt;/td&gt;
&lt;td&gt;6.367 s&lt;/td&gt;
&lt;td&gt;1.328 s&lt;/td&gt;
&lt;td&gt;1.054 s&lt;/td&gt;
&lt;td&gt;8.62 GB&lt;/td&gt;
&lt;td&gt;Leverages repeating 1D spatial structures.&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;&lt;strong&gt;Spatial Transfer Matrix (B5)&lt;/strong&gt;&lt;/td&gt;
&lt;td&gt;7.375 s&lt;/td&gt;
&lt;td&gt;1.320 s&lt;/td&gt;
&lt;td&gt;1.133 s&lt;/td&gt;
&lt;td&gt;5.44 GB&lt;/td&gt;
&lt;td&gt;A lower-VRAM variant of the same exact STM algorithm.&lt;/td&gt;
&lt;/tr&gt;
&lt;/tbody&gt;
&lt;/table&gt;&lt;/div&gt;

&lt;p&gt;These approaches complement one another:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;
&lt;strong&gt;The Global Tree&lt;/strong&gt; delivers excellent steady-state execution time (1.09s) but suffers from a brutal 5+ hour cold compilation time and massive 52 GB peak VRAM usage.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;The Local Window&lt;/strong&gt; approach crushes compile time down to ~10 seconds, but redundantly computes the shared environments of adjacent local terms, making single-GPU execution sluggish.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;The Spatial Transfer Matrix (STM)&lt;/strong&gt; successfully decouples compilation costs from the system length by explicitly passing the repetitive 1D spatial structure to &lt;code&gt;jax.lax.scan&lt;/code&gt;, perfectly balancing compilation and execution efficiency.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;&lt;a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fa4ko4snbtifzomsp4l11.png" class="article-body-image-wrapper"&gt;&lt;img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fa4ko4snbtifzomsp4l11.png" alt="scatter plot for efficiency of three methods" width="800" height="591"&gt;&lt;/a&gt;&lt;/p&gt;

&lt;h2&gt;
  
  
  Contraction Path Metrics: FLOPs, Max Tensor, and Write
&lt;/h2&gt;

&lt;p&gt;In tensor network contraction, path quality is defined by several resource metrics:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;
&lt;strong&gt;FLOPs (Floating Point Operations)&lt;/strong&gt;: The total number of scalar multiply-add operations required for the path. For example, &lt;code&gt;log10 FLOPs = 12.0&lt;/code&gt; means roughly $10^{12}$ scalar ops.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;Max Tensor Size (Width $w$)&lt;/strong&gt;: The number of elements in the largest intermediate tensor generated during contraction. It directly dictates the instantaneous VRAM pressure during the forward pass. (If all indices have a dimension of 2, width $w$ corresponds to $2^w$ elements).&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;Total Write&lt;/strong&gt;: The total number of elements written out across all intermediate tensors. During backpropagation, &lt;em&gt;all&lt;/em&gt; of these written tensors must be retained in memory as residuals to compute gradients.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;When calculating expectations and gradients simultaneously, you must keep an eye on all three metrics. Ultimately, however, empirical compile time, execution time, and peak GPU VRAM are the ground truth.&lt;/p&gt;

&lt;h2&gt;
  
  
  Perspective 1: Global Contraction Tree — Brute-Forcing the Entire Graph
&lt;/h2&gt;

&lt;p&gt;By calling &lt;code&gt;value_and_grad&lt;/code&gt; directly on the complete bra-MPO-ket network, we can quantify the true cost of naive global contraction.&lt;/p&gt;

&lt;p&gt;To make this feasible, we first alter the network representation. It is critical to write each $R_{ZZ}$ gate in an exact Rank-2 decomposed form:&lt;/p&gt;

&lt;p&gt;$$R_{ZZ}(\theta)=\cos(\theta/2)I\otimes I-i\sin(\theta/2)Z\otimes Z$$&lt;/p&gt;

&lt;p&gt;Maintaining the original left-to-right "ladder" gate ordering is also crucial. If we reorder the commuting $R_{ZZ}$ gates into an alternating even/odd bond "brick-wall" structure, we destroy the elimination path. The tensor network width instantly spikes to 39.585. Keeping the ladder ordering keeps the width at a manageable 23.585.&lt;/p&gt;

&lt;p&gt;Using the &lt;code&gt;omeco&lt;/code&gt; path search on this massive graph, we found a global contraction path with &lt;code&gt;log10 FLOPs = 12.0466&lt;/code&gt; and &lt;code&gt;log2 write = 32.5912&lt;/code&gt;.&lt;/p&gt;

&lt;p&gt;On a single RTX 6000D, post-compilation gradient execution takes only 1.09s, peaking at 51 GB of VRAM. This proves that full-graph VQE gradients can technically be run on a single card, but the bottlenecks are purely the global HLO compilation time and the memory needed to store backward residuals.&lt;/p&gt;

&lt;h2&gt;
  
  
  omeco vs. cotengra: Two Path-Search Frameworks
&lt;/h2&gt;

&lt;p&gt;Since the global graph's resource bottleneck lies heavily in contraction path search and compilation, how do top-tier tools fare? We compared two frameworks—&lt;code&gt;omeco&lt;/code&gt; and &lt;code&gt;cotengra&lt;/code&gt;—on this massive tensor network.&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;
&lt;strong&gt;omeco&lt;/strong&gt; (Rust-based) uses a TreeSA algorithm to perform intense simulated annealing in tree space. It is excellent for quickly finding high-quality un-sliced seed trees on large graphs.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;cotengra&lt;/strong&gt; offers richer interfaces for manipulating and fine-tuning contraction trees, such as reconfiguring subtrees based on existing paths or performing granular slice-index searches on a fixed tree.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;We ran a controlled search comparison on a 100-qubit, 10-layer TFIM graph (6,460 tensors, 8,441 indices):&lt;/p&gt;

&lt;div class="table-wrapper-paragraph"&gt;&lt;table&gt;
&lt;thead&gt;
&lt;tr&gt;
&lt;th&gt;&lt;strong&gt;Search Method &amp;amp; Budget&lt;/strong&gt;&lt;/th&gt;
&lt;th&gt;&lt;strong&gt;Search Time&lt;/strong&gt;&lt;/th&gt;
&lt;th&gt;&lt;strong&gt;log10 FLOPs&lt;/strong&gt;&lt;/th&gt;
&lt;th&gt;&lt;strong&gt;log2 Max Tensor&lt;/strong&gt;&lt;/th&gt;
&lt;th&gt;&lt;strong&gt;log2 Total Write&lt;/strong&gt;&lt;/th&gt;
&lt;th&gt;&lt;strong&gt;Notes&lt;/strong&gt;&lt;/th&gt;
&lt;/tr&gt;
&lt;/thead&gt;
&lt;tbody&gt;
&lt;tr&gt;
&lt;td&gt;
&lt;code&gt;omeco&lt;/code&gt; TreeSA (16 trials × 64 steps)&lt;/td&gt;
&lt;td&gt;15.68 s&lt;/td&gt;
&lt;td&gt;11.068&lt;/td&gt;
&lt;td&gt;24.585&lt;/td&gt;
&lt;td&gt;29.517&lt;/td&gt;
&lt;td&gt;16 independent SA trials&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;
&lt;code&gt;cotengra&lt;/code&gt; CMA-ES (8 hyperparam candidates)&lt;/td&gt;
&lt;td&gt;37.59 s&lt;/td&gt;
&lt;td&gt;12.901&lt;/td&gt;
&lt;td&gt;31.585&lt;/td&gt;
&lt;td&gt;35.646&lt;/td&gt;
&lt;td&gt;8 sets of search params&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;
&lt;code&gt;cotengra&lt;/code&gt; SA (12 steps × 12 reconfigs)&lt;/td&gt;
&lt;td&gt;9.22 s&lt;/td&gt;
&lt;td&gt;11.827&lt;/td&gt;
&lt;td&gt;29.585&lt;/td&gt;
&lt;td&gt;34.014&lt;/td&gt;
&lt;td&gt;Low-budget SA&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;
&lt;code&gt;cotengra&lt;/code&gt; SA (24 steps × 24 reconfigs)&lt;/td&gt;
&lt;td&gt;34.09 s&lt;/td&gt;
&lt;td&gt;10.873&lt;/td&gt;
&lt;td&gt;23.585&lt;/td&gt;
&lt;td&gt;30.367&lt;/td&gt;
&lt;td&gt;Higher-budget SA&lt;/td&gt;
&lt;/tr&gt;
&lt;/tbody&gt;
&lt;/table&gt;&lt;/div&gt;

&lt;p&gt;The data shows that &lt;code&gt;omeco&lt;/code&gt; TreeSA ran significantly more annealing steps in less time and returned lower-complexity trees, outperforming the alternatives for this un-sliced scenario. However, for scenarios requiring deep slicing, neither framework is currently ideal.&lt;/p&gt;

&lt;h2&gt;
  
  
  The Efficiency Limits of Slicing
&lt;/h2&gt;

&lt;p&gt;When VRAM is tight, slicing the contraction tree is a necessary evil to keep memory footprints in check. But if you have enough VRAM for an un-sliced path, forcing a slice to distribute across multiple GPUs does not scale efficiency linearly.&lt;/p&gt;

&lt;p&gt;Fundamentally, slicing a tensor network means forcibly delaying the contraction of selected indices to the very end of the path. In an optimal, un-sliced path, these indices are eliminated early on, preventing intermediate tensors from blowing up in size. Slicing disrupts this optimal elimination order. It forces every sub-task to redundantly compute intermediate results that could have been shared before the final summation.&lt;/p&gt;

&lt;p&gt;When we force-sliced the 1,000-qubit global path into 8, 32, and 128 tasks, the &lt;code&gt;log2 write&lt;/code&gt; per slice barely dropped—from 32.5912 to 32.5632, 32.5491, and 32.5365 respectively. Because it failed to meaningfully dismantle the primary memory structure of the backward pass, while heavily deviating from the optimal contraction order, redundant computation skyrocketed. Total &lt;code&gt;log10 FLOPs&lt;/code&gt; jumped from 12.0466 to 12.9420, 13.5391, and 14.1355.&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;Takeaway:&lt;/strong&gt; If it fits in VRAM, using an un-sliced path usually maximizes hardware utilization.&lt;/p&gt;

&lt;p&gt;&lt;a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2F5o3kezbmdg3v1djilwap.png" class="article-body-image-wrapper"&gt;&lt;img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2F5o3kezbmdg3v1djilwap.png" alt="framework of the three methods" width="799" height="270"&gt;&lt;/a&gt;&lt;/p&gt;

&lt;h2&gt;
  
  
  Perspective 2: Local Causal Sliding Window — Exploiting Local Physics
&lt;/h2&gt;

&lt;p&gt;For a circuit depth of $L=10$, we can leverage the commutativity of 2-qubit gates to convert the ladder structure into a brick-wall equivalent. Consequently, every local TFIM term possesses an exact backward causal cone during backpropagation. Since the diagonal entangling gates commute, the width of this causal window is strictly $2L+2$, which is only 22 sites for $L=10$. The brick-wall structure, which was a liability for global tree search, becomes a massive advantage here.&lt;/p&gt;

&lt;p&gt;By using &lt;code&gt;lax.scan&lt;/code&gt; to iterate over the local Pauli terms of the Hamiltonian and wrapping single terms in &lt;code&gt;jax.checkpoint&lt;/code&gt;, the compilation scale becomes solely dependent on window size and circuit depth, totally decoupled from the total number of system sites:&lt;/p&gt;

&lt;p&gt;$$T_{\rm window}=O(n C_{\rm local}(L)),\qquad M_{\rm window}=O(M_{\rm local}(L))$$&lt;/p&gt;

&lt;p&gt;This approach boasts rapid compile times (~10 seconds) and a lean VRAM footprint of 8.86 GB. If dispatched across 8 GPUs (125 local terms per GPU), the steady-state gradient time drops to 18.71s. However, the overlapping causal windows introduce massive computational redundancy.&lt;/p&gt;

&lt;p&gt;Furthermore, this method heavily relies on the commutativity of diagonal gates. If we swap them for non-commuting 2-qubit gates, the causal cone under a ladder structure expands to $O(n)$, rendering this method useless.&lt;/p&gt;

&lt;h2&gt;
  
  
  Perspective 3: Spatial Transfer Matrix — Lossless Exact Scanning
&lt;/h2&gt;

&lt;p&gt;The spatial transfer matrix strategy slices the system along the spatial qubit dimension into a left-end block, repeating middle blocks, and a right-end block. The middle block receives the complete boundary tensor from the left, contracts its internal bra-MPO-ket subnetwork, and outputs to the next block.&lt;/p&gt;

&lt;p&gt;The recurrence relation is:&lt;/p&gt;

&lt;p&gt;$$B_{k+1}=\mathcal{T}_k(\boldsymbol\theta_k)B_k$$&lt;/p&gt;

&lt;p&gt;At depth $L=10$, the spatial cut crosses 10 ket bonds, 10 bra bonds, and 1 MPO bond (dimension 3). The exact boundary size is:&lt;/p&gt;

&lt;p&gt;$$D_{\rm boundary}=3\times2^{2L}=3\times4^L = 3\times4^{10}$$&lt;/p&gt;

&lt;p&gt;This amounts to roughly 24 MiB of memory traffic, which perfectly explains why this method can effortlessly perform linear scanning at the 1,000-qubit scale.&lt;/p&gt;

&lt;p&gt;Using a 3-site block division, this mode requires just 6.367s for cold compilation, executes at 1.054s per step, and peaks at 8.62 GB of VRAM.&lt;/p&gt;

&lt;p&gt;The spatial width of each scan step (the block size $b$) acts as a tunable time-space tradeoff:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;
&lt;strong&gt;Smaller blocks&lt;/strong&gt; shrink the internal contraction graph, reducing compile pressure and the residuals saved for backprop, but increase the number of scan steps needed to pass the boundary tensor.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;Larger blocks&lt;/strong&gt; reduce scan steps but inflate maximum tensor size, total write, compile time, and peak VRAM.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;In our tests, a 3-site block offered optimal throughput. A 5-site block reduced peak VRAM further to 5.44 GB, at the slight cost of increasing execution time to 1.133s.&lt;/p&gt;

&lt;h2&gt;
  
  
  Counterintuitive Acceleration: Trading Compute for Bandwidth
&lt;/h2&gt;

&lt;p&gt;Fascinatingly, pairing the transfer matrix's middle block loop with &lt;code&gt;jax.checkpoint&lt;/code&gt; (recomputation) dropped the peak VRAM from 42.24 GB down to 8.62 GB, while &lt;em&gt;also&lt;/em&gt; slightly reducing steady-state execution time from 1.118s to 1.054s.&lt;/p&gt;

&lt;p&gt;"Reducing memory traffic by recomputing intermediate residuals" actually speeding up execution highlights a hardware reality: on modern consumer GPUs, memory bandwidth is often a much harder bottleneck than raw single-precision FLOPs. Taking on a heavier compute burden to alleviate read/write strain yields net performance gains. This same logic applies when tuning the penalty ratio between time and space complexity during contraction tree searches.&lt;/p&gt;

&lt;p&gt;Even more thought-provoking is the fact that this relatively simple spatial slicing and block-scanning strategy heavily outperformed the global tree—derived from intense simulated annealing—in execution time, compile time, and memory footprint. This implies that for truly massive tensor networks requiring backpropagation (where intermediate residuals and VRAM are paramount), existing contraction path algorithms still have significant blind spots and ample room for innovation.&lt;/p&gt;

&lt;h2&gt;
  
  
  Pitfalls to Avoid
&lt;/h2&gt;

&lt;ol&gt;
&lt;li&gt;&lt;strong&gt;Bypassing TF32 Precision Limitations&lt;/strong&gt;&lt;/li&gt;
&lt;/ol&gt;

&lt;p&gt;When doing complex matrix multiplications on NVIDIA GPUs, simply setting the highest precision in JAX is not enough to avoid automatic degradation to TF32. You must explicitly set this environment variable before runtime:&lt;/p&gt;

&lt;p&gt;Bash&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight plaintext"&gt;&lt;code&gt;   export NVIDIA_TF32_OVERRIDE=0
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;If you don't, the state norm can slip to 0.98–0.996, causing noticeable deviations in energy calculations. This is a stark reminder of the differing precision tolerances between standard Machine Learning and Quantum Simulation.&lt;/p&gt;

&lt;ol&gt;
&lt;li&gt;&lt;strong&gt;Rayon Thread Pool Stack Overflows&lt;/strong&gt;&lt;/li&gt;
&lt;/ol&gt;

&lt;p&gt;When &lt;code&gt;omeco&lt;/code&gt; searches ultra-deep trees (depth &amp;gt; 1000) and generates slicing code, the recursive traversal of the deep tree can trigger a Segmentation Fault in Rayon threads. You must explicitly bump up the stack size before initiating Python or the Rayon global pool:&lt;/p&gt;

&lt;p&gt;Bash&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight plaintext"&gt;&lt;code&gt;   export RUST_MIN_STACK=67108864
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;h2&gt;
  
  
  Full 1,000-Qubit VQE: From Single-Step Gradients to End-to-End Optimization
&lt;/h2&gt;

&lt;p&gt;To prove that this isn't just a parlor trick for a single energy/gradient calculation, we ran a complete, complex 1,000-qubit VQE optimization. Our circuit features 10 layers of nearest-neighbor $R_{ZZ}$ ladders, with each layer's single-qubit component consisting of sequential $R_X$, $R_Z$, and $R_X$ rotations:&lt;/p&gt;

&lt;p&gt;$$\vert{}0\rangle^{\otimes 1000}\xrightarrow{H^{\otimes 1000}}\left[R_{ZZ}\text{ Ladder}\rightarrow R_X\rightarrow R_Z\rightarrow R_X\right]^{10}.$$&lt;/p&gt;

&lt;p&gt;Before constructing the full bra-MPO-ket network, we exactly compressed each consecutive $R_X$–$R_Z$–$R_X$ sequence into a single $2\times2$ single-qubit tensor. We then utilized the exact same Spatial Transfer Matrix scheme (with zero approximation truncations) to compute the energy and the gradients for all 40,000 parameters. The entire optimization ran on a single NVIDIA RTX 6000D, completing 2,000 Adam updates.&lt;/p&gt;

&lt;p&gt;For this 1,000-site critical open-boundary TFIM, the exact ground state energy can be derived directly from the free-fermion spectrum. After 2,000 updates, our VQE energy was $E_{\rm VQE}=-1268.4365$, resulting in a total energy error of $E_{\rm VQE}-E_0=4.440$. This equates to a relative error of about $3\times10^{-3}$.&lt;/p&gt;

&lt;p&gt;Without any hyperparameter tuning, a 1,000-qubit VQE featuring complex gate sequences, full parameters, and exact full gradients hit a relative energy precision in the thousandths. This confirms that our tensor network representations, path searches, spatial scanning, autodiff, and GPU execution aren't isolated micro-optimizations—they form a cohesive, end-to-end technical pipeline capable of translating circuit definitions into actual, strictly-verified optimization convergence.&lt;/p&gt;

&lt;h2&gt;
  
  
  Conclusion
&lt;/h2&gt;

&lt;p&gt;Our work on full-gradient simulation and optimization for a 1,000-qubit VQE demonstrates that by hybridizing low-rank gate decomposition, contraction path search, lossless spatial transfer matrices, JAX automatic differentiation, and compiler mechanics, we can do the impossible. We can compute the full gradient of a 1,000-qubit VQE rapidly on a single GPU, execute thousands of training steps within hours, and achieve ~0.3% relative energy error without systemic fine-tuning.&lt;/p&gt;

&lt;p&gt;These results illustrate the immense potential of moving from single-shot simulations to massive end-to-end variational computation, leaving plenty of headroom for exploring better circuit structures, learning rate schedules, and higher precision.&lt;/p&gt;

&lt;p&gt;Getting a 1,000-qubit VQE to run natively proves one thing: in the deep waters of quantum simulation, the mathematical elegance of algorithmic design must be deeply coupled with the underlying mechanics of compilers (XLA), automatic differentiation, and GPU hardware architectures. TensorCircuit-NG is more than just a simulator; it is a critical piece of infrastructure that bridges software and hardware, providing a complete technical pipeline from quantum algorithm definition to ultra-large-scale variational optimization.&lt;/p&gt;

</description>
      <category>ai</category>
      <category>jax</category>
      <category>quantum</category>
    </item>
    <item>
      <title>Large-Scale TensorCircuit Contractions on GPUs: Disabling XLA GPU Autotuning</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Wed, 15 Jul 2026 03:41:23 +0000</pubDate>
      <link>https://dev.to/refractionray/large-scale-tensorcircuit-contractions-on-gpus-disabling-xla-gpu-autotuning-3p2</link>
      <guid>https://dev.to/refractionray/large-scale-tensorcircuit-contractions-on-gpus-disabling-xla-gpu-autotuning-3p2</guid>
      <description>&lt;p&gt;When running large-scale tensor-network contractions with TensorCircuit-NG and the JAX GPU backend, the following runtime configuration is worth testing:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight shell"&gt;&lt;code&gt;&lt;span class="nv"&gt;XLA_PYTHON_CLIENT_PREALLOCATE&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="nb"&gt;false &lt;/span&gt;&lt;span class="nv"&gt;XLA_FLAGS&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="nt"&gt;--xla_gpu_autotune_level&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;0 python your_script.py
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;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.&lt;/p&gt;

&lt;h2&gt;
  
  
  Key takeaway
&lt;/h2&gt;

&lt;p&gt;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.&lt;/p&gt;

&lt;p&gt;For TensorCircuit contractions, run this A/B test by default:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight shell"&gt;&lt;code&gt;&lt;span class="c"&gt;# Baseline&lt;/span&gt;
&lt;span class="nv"&gt;XLA_PYTHON_CLIENT_PREALLOCATE&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="nb"&gt;false &lt;/span&gt;python your_script.py

&lt;span class="c"&gt;# Test configuration&lt;/span&gt;
&lt;span class="nv"&gt;XLA_PYTHON_CLIENT_PREALLOCATE&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="nb"&gt;false &lt;/span&gt;&lt;span class="nv"&gt;XLA_FLAGS&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="nt"&gt;--xla_gpu_autotune_level&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;0 python your_script.py
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;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.&lt;/p&gt;

&lt;h2&gt;
  
  
  Representative results
&lt;/h2&gt;

&lt;p&gt;All results below use a fixed contraction path so that path-search randomness does not affect the comparison.&lt;/p&gt;

&lt;div class="table-wrapper-paragraph"&gt;&lt;table&gt;
&lt;thead&gt;
&lt;tr&gt;
&lt;th&gt;Workload&lt;/th&gt;
&lt;th&gt;Autotuning&lt;/th&gt;
&lt;th&gt;Post-compile memory&lt;/th&gt;
&lt;th&gt;Peak memory&lt;/th&gt;
&lt;th&gt;Steady-state runtime&lt;/th&gt;
&lt;/tr&gt;
&lt;/thead&gt;
&lt;tbody&gt;
&lt;tr&gt;
&lt;td&gt;100 qubits × 24 layers, amplitude&lt;/td&gt;
&lt;td&gt;Default&lt;/td&gt;
&lt;td&gt;8.7 GiB&lt;/td&gt;
&lt;td&gt;8.7 GiB&lt;/td&gt;
&lt;td&gt;0.37 s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;100 qubits × 24 layers, amplitude&lt;/td&gt;
&lt;td&gt;autotune=0&lt;/td&gt;
&lt;td&gt;0.5 GiB&lt;/td&gt;
&lt;td&gt;4.6 GiB&lt;/td&gt;
&lt;td&gt;0.40 s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;28 qubits × 12 layers, expectation gradient&lt;/td&gt;
&lt;td&gt;Default&lt;/td&gt;
&lt;td&gt;4.7 GiB&lt;/td&gt;
&lt;td&gt;21.4 GiB&lt;/td&gt;
&lt;td&gt;0.18 s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;28 qubits × 12 layers, expectation gradient&lt;/td&gt;
&lt;td&gt;autotune=0&lt;/td&gt;
&lt;td&gt;0.5 GiB&lt;/td&gt;
&lt;td&gt;17.2 GiB&lt;/td&gt;
&lt;td&gt;0.15 s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;28 qubits × 13 layers, expectation gradient&lt;/td&gt;
&lt;td&gt;Default&lt;/td&gt;
&lt;td&gt;10.9 GiB&lt;/td&gt;
&lt;td&gt;43.8 GiB&lt;/td&gt;
&lt;td&gt;0.46 s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;28 qubits × 13 layers, expectation gradient&lt;/td&gt;
&lt;td&gt;autotune=0&lt;/td&gt;
&lt;td&gt;0.5 GiB&lt;/td&gt;
&lt;td&gt;33.5 GiB&lt;/td&gt;
&lt;td&gt;0.42 s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;28 qubits × 14 layers, expectation gradient&lt;/td&gt;
&lt;td&gt;Default&lt;/td&gt;
&lt;td&gt;17.0 GiB&lt;/td&gt;
&lt;td&gt;64.5 GiB&lt;/td&gt;
&lt;td&gt;0.58 s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;28 qubits × 14 layers, expectation gradient&lt;/td&gt;
&lt;td&gt;autotune=0&lt;/td&gt;
&lt;td&gt;0.5 GiB&lt;/td&gt;
&lt;td&gt;64.5 GiB&lt;/td&gt;
&lt;td&gt;0.58 s&lt;/td&gt;
&lt;/tr&gt;
&lt;/tbody&gt;
&lt;/table&gt;&lt;/div&gt;

&lt;p&gt;The first three workload groups show the most common benefit: disabling autotuning reduces compilation-stage memory and the final peak. The 28 × 14 case appears different: post-compile memory falls from 17.0 GiB to 0.5 GiB, but the final peak is unchanged. This behavior is related to the JAX GPU allocator.&lt;/p&gt;

&lt;h2&gt;
  
  
  Why the final peak is unchanged for the 28 × 14 case
&lt;/h2&gt;

&lt;p&gt;Even with &lt;code&gt;XLA_PYTHON_CLIENT_PREALLOCATE=false&lt;/code&gt;, the default JAX GPU allocator tends to retain and reuse GPU memory that it has already allocated. The memory reported by &lt;code&gt;nvidia-smi&lt;/code&gt; is therefore the amount held by the process, not the total size of tensors that are currently live.&lt;/p&gt;

&lt;p&gt;In the 28 × 14 example, default autotuning did consume approximately 17 GiB of additional memory during compilation. During the first real backward contraction, however, the runtime buffers themselves also required a large allocation. The default allocator could reuse blocks allocated earlier, so the final peak was not simply the sum of runtime memory and autotuning memory.&lt;/p&gt;

&lt;p&gt;Using &lt;code&gt;XLA_PYTHON_CLIENT_ALLOCATOR=platform&lt;/code&gt; as a diagnostic exposes this difference:&lt;/p&gt;

&lt;div class="table-wrapper-paragraph"&gt;&lt;table&gt;
&lt;thead&gt;
&lt;tr&gt;
&lt;th&gt;Workload&lt;/th&gt;
&lt;th&gt;Allocator&lt;/th&gt;
&lt;th&gt;Autotuning&lt;/th&gt;
&lt;th&gt;Post-compile memory&lt;/th&gt;
&lt;th&gt;Peak memory&lt;/th&gt;
&lt;/tr&gt;
&lt;/thead&gt;
&lt;tbody&gt;
&lt;tr&gt;
&lt;td&gt;28 qubits × 14 layers, expectation gradient&lt;/td&gt;
&lt;td&gt;Default&lt;/td&gt;
&lt;td&gt;Default&lt;/td&gt;
&lt;td&gt;17.0 GiB&lt;/td&gt;
&lt;td&gt;64.5 GiB&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;28 qubits × 14 layers, expectation gradient&lt;/td&gt;
&lt;td&gt;Default&lt;/td&gt;
&lt;td&gt;autotune=0&lt;/td&gt;
&lt;td&gt;0.5 GiB&lt;/td&gt;
&lt;td&gt;64.5 GiB&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;28 qubits × 14 layers, expectation gradient&lt;/td&gt;
&lt;td&gt;platform&lt;/td&gt;
&lt;td&gt;Default&lt;/td&gt;
&lt;td&gt;17.0 GiB&lt;/td&gt;
&lt;td&gt;82.8 GiB&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;28 qubits × 14 layers, expectation gradient&lt;/td&gt;
&lt;td&gt;platform&lt;/td&gt;
&lt;td&gt;autotune=0&lt;/td&gt;
&lt;td&gt;0.5 GiB&lt;/td&gt;
&lt;td&gt;66.3 GiB&lt;/td&gt;
&lt;/tr&gt;
&lt;/tbody&gt;
&lt;/table&gt;&lt;/div&gt;

&lt;p&gt;The &lt;code&gt;platform&lt;/code&gt; allocator retains less reusable GPU memory, which makes it useful for diagnosing where allocations originate. Because it also reuses fewer large blocks allocated earlier, however, the runtime may request additional memory and produce a higher peak. Prefer the default allocator for normal execution; use &lt;code&gt;platform&lt;/code&gt; mainly to confirm whether autotuning introduces extra memory usage.&lt;/p&gt;

&lt;h2&gt;
  
  
  Interpreting peak memory
&lt;/h2&gt;

&lt;p&gt;Here, peak memory is the maximum per-process GPU memory observed by polling &lt;code&gt;nvidia-smi&lt;/code&gt;, with snapshots recorded at stages such as &lt;code&gt;after_compile&lt;/code&gt; and &lt;code&gt;after_first_run&lt;/code&gt;. XLA's &lt;code&gt;compiled.memory_analysis().peak_memory_in_bytes&lt;/code&gt; measures something different: it more closely reflects the computation graph's buffer assignment and is usually lower than the actual process memory reported by &lt;code&gt;nvidia-smi&lt;/code&gt;.&lt;/p&gt;

&lt;h2&gt;
  
  
  Why the first GPU may use more memory in multi-GPU processes
&lt;/h2&gt;

&lt;p&gt;In a JAX/XLA process that can see multiple GPUs, backend initialization, the default device, compilation services, autotuning caches, executable caches, or allocator state may be placed preferentially on the first GPU in the visible-device list. Consequently, the first GPU can hold an extra block of memory. Disabling &lt;code&gt;xla_gpu_autotune_level&lt;/code&gt; often reduces this first-device memory tax.&lt;/p&gt;

&lt;p&gt;For multiple independent single-GPU jobs, use &lt;code&gt;CUDA_VISIBLE_DEVICES&lt;/code&gt; so that each process sees only its assigned GPU:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight shell"&gt;&lt;code&gt;&lt;span class="nv"&gt;CUDA_VISIBLE_DEVICES&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;1 &lt;span class="nv"&gt;XLA_PYTHON_CLIENT_PREALLOCATE&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="nb"&gt;false &lt;/span&gt;&lt;span class="nv"&gt;XLA_FLAGS&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="nt"&gt;--xla_gpu_autotune_level&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;0 python your_script.py
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;For genuine multi-GPU parallel or distributed jobs, do not hide GPUs that must participate in the computation merely to avoid the first-device memory tax. Preserve the correct device set, then measure how disabling autotuning changes memory usage on each GPU.&lt;/p&gt;

&lt;h2&gt;
  
  
  Practical recommendations
&lt;/h2&gt;

&lt;p&gt;For large TensorCircuit-NG contractions, use the following procedure:&lt;/p&gt;

&lt;ol&gt;
&lt;li&gt;Fix the contraction path instead of running a new stochastic path search for every benchmark.&lt;/li&gt;
&lt;li&gt;Set &lt;code&gt;XLA_PYTHON_CLIENT_PREALLOCATE=false&lt;/code&gt; to prevent JAX from reserving a large block of GPU memory at startup.&lt;/li&gt;
&lt;li&gt;A/B test default autotuning against &lt;code&gt;XLA_FLAGS=--xla_gpu_autotune_level=0&lt;/code&gt;.&lt;/li&gt;
&lt;li&gt;Record compile time, first-run time, steady-state runtime, post-compile memory, and post-first-run memory separately.&lt;/li&gt;
&lt;li&gt;Even if disabling autotuning does not improve steady-state runtime, prefer &lt;code&gt;autotune=0&lt;/code&gt;, especially near the OOM limit or when scheduling multiple GPUs.&lt;/li&gt;
&lt;/ol&gt;

&lt;p&gt;In short, disabling XLA GPU autotuning does not change the contraction path; it reduces an additional source of memory usage in the GPU compilation and execution layers. For large TensorCircuit contractions, this is usually a low-risk configuration well worth testing.&lt;/p&gt;

</description>
      <category>jax</category>
    </item>
    <item>
      <title>When Quantum Dynamics Doesn't Start from Zero: Completing the Missing Half of Entanglement Growth | PRL</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Fri, 10 Jul 2026 00:42:59 +0000</pubDate>
      <link>https://dev.to/refractionray/when-quantum-dynamics-doesnt-start-from-zero-completing-the-missing-half-of-entanglement-growth--pl7</link>
      <guid>https://dev.to/refractionray/when-quantum-dynamics-doesnt-start-from-zero-completing-the-missing-half-of-entanglement-growth--pl7</guid>
      <description>&lt;p&gt;Entanglement dynamics lies at the heart of nonequilibrium quantum physics. For more than two decades, the standard approach has been remarkably consistent: start from an unentangled product state, let the system evolve unitarily, and study how the entanglement entropy grows over time.&lt;/p&gt;

&lt;p&gt;This framework has led to many fundamental discoveries, including our understanding of quantum thermalization and many-body localization (MBL). It also shaped an implicit assumption shared across the field: whenever the half-chain entanglement entropy increases, the system must be &lt;em&gt;creating&lt;/em&gt; new quantum entanglement.&lt;/p&gt;

&lt;p&gt;Our recent paper, published in &lt;em&gt;Physical Review Letters&lt;/em&gt;, argues that this picture is incomplete.&lt;/p&gt;

&lt;p&gt;The central observation is surprisingly simple:&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;&lt;strong&gt;An increase in measured entanglement does not necessarily mean new entanglement has been created.&lt;/strong&gt;&lt;/p&gt;
&lt;/blockquote&gt;

&lt;p&gt;Instead, quantum dynamics possesses two fundamentally different capabilities that have long been mixed together.&lt;/p&gt;




&lt;h2&gt;
  
  
  Two ways to increase entanglement
&lt;/h2&gt;

&lt;p&gt;A useful way to think about entanglement is to imagine the system as a connected network of water reservoirs.&lt;/p&gt;

&lt;p&gt;There are two distinct mechanisms that can raise the water level observed at a particular cut.&lt;/p&gt;

&lt;h3&gt;
  
  
  Build: Creating new entanglement
&lt;/h3&gt;

&lt;p&gt;The first mechanism genuinely generates new quantum correlations.&lt;/p&gt;

&lt;p&gt;This is analogous to pumping fresh water into the entire reservoir system. The total amount of entanglement increases.&lt;/p&gt;

&lt;h3&gt;
  
  
  Move: Transporting existing entanglement
&lt;/h3&gt;

&lt;p&gt;The second mechanism creates &lt;strong&gt;no new entanglement at all&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;Instead, quantum evolution simply redistributes the entanglement that already exists inside the system. Water is flowing internally, but no new water enters the reservoirs. Although the total amount remains unchanged, the water level measured at one particular location may still rise.&lt;/p&gt;

&lt;p&gt;Exactly the same phenomenon can happen for half-chain entanglement entropy.&lt;/p&gt;

&lt;p&gt;This distinction turns out to be much more important than previously appreciated.&lt;/p&gt;




&lt;h2&gt;
  
  
  The intuitive expectation
&lt;/h2&gt;

&lt;p&gt;Once these two mechanisms are separated, an intuitive principle emerges.&lt;/p&gt;

&lt;p&gt;If a system already starts with a large amount of entanglement, there is less room left to generate additional entanglement later.&lt;/p&gt;

&lt;p&gt;Indeed, this is exactly what happens in familiar systems such as&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;quantum chaotic (thermalizing) systems,&lt;/li&gt;
&lt;li&gt;free-fermion systems,&lt;/li&gt;
&lt;li&gt;random quantum circuits with strong scrambling.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;As the initial entanglement increases, the additional entanglement generated during evolution decreases monotonically.&lt;/p&gt;

&lt;p&gt;This is precisely what one would expect if entanglement growth were primarily driven by &lt;strong&gt;building&lt;/strong&gt; new entanglement.&lt;/p&gt;




&lt;h2&gt;
  
  
  MBL breaks the rule
&lt;/h2&gt;

&lt;p&gt;Many-body localized systems tell a completely different story.&lt;/p&gt;

&lt;p&gt;Instead of decreasing monotonically, the entanglement growth exhibits a striking &lt;strong&gt;bell-shaped dependence&lt;/strong&gt; on the initial entanglement.&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;Starting from nearly product states, the growth is very small.&lt;/li&gt;
&lt;li&gt;Starting from nearly maximally entangled states, the growth is again very small.&lt;/li&gt;
&lt;li&gt;The largest increase occurs at &lt;strong&gt;intermediate initial entanglement&lt;/strong&gt;.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;This behavior is difficult to explain if entanglement growth comes solely from generating new entanglement.&lt;/p&gt;

&lt;p&gt;Something else must be happening.&lt;/p&gt;




&lt;h2&gt;
  
  
  A "pure transport" experiment
&lt;/h2&gt;

&lt;p&gt;To isolate the missing ingredient, we designed a particularly simple reference model.&lt;/p&gt;

&lt;p&gt;Instead of using an interacting Hamiltonian, we considered a random circuit composed &lt;strong&gt;only of SWAP gates&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;This circuit has a remarkable property:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;it &lt;strong&gt;cannot create any entanglement&lt;/strong&gt;, and&lt;/li&gt;
&lt;li&gt;it only exchanges the locations of quantum states.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;In other words, it performs &lt;strong&gt;pure entanglement transport&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;Surprisingly, the entanglement-growth curve produced by this SWAP-only circuit closely matches the behavior observed in MBL systems, both qualitatively and even quantitatively.&lt;/p&gt;

&lt;p&gt;This comparison reveals the underlying physics.&lt;/p&gt;

&lt;p&gt;MBL is not particularly good at generating new entanglement.&lt;/p&gt;

&lt;p&gt;Instead, it excels at &lt;strong&gt;moving around the entanglement that already exists&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;When the system begins with almost no entanglement, there is simply nothing to transport.&lt;/p&gt;

&lt;p&gt;When it begins nearly saturated, there is no remaining room for rearrangement.&lt;/p&gt;

&lt;p&gt;Only at intermediate entanglement do both ingredients coexist, allowing transport to produce the largest observable increase.&lt;/p&gt;




&lt;h2&gt;
  
  
  Looking beyond a single bipartition
&lt;/h2&gt;

&lt;p&gt;Most previous studies monitor only one quantity: the entanglement across the middle cut of the system.&lt;/p&gt;

&lt;p&gt;While extremely useful, this provides only a partial view of the system's entanglement structure.&lt;/p&gt;

&lt;p&gt;In our work, we introduce a more global perspective by considering the &lt;strong&gt;Bipartition-Averaged Entanglement Entropy (BAEE)&lt;/strong&gt;, which averages entanglement over all possible bipartitions.&lt;/p&gt;

&lt;p&gt;The simulations reveal an interesting phenomenon.&lt;/p&gt;

&lt;p&gt;Even in ordinary thermalizing systems, BAEE grows much faster than the half-chain entanglement during the early stages of evolution.&lt;/p&gt;

&lt;p&gt;The difference between these two quantities represents a hidden reservoir of entanglement that has already been generated locally but has not yet reached the particular cut being measured.&lt;/p&gt;

&lt;p&gt;Returning to our water analogy, thermalization rapidly fills many local reservoirs throughout the system.&lt;/p&gt;

&lt;p&gt;Later dynamics can transport this stored entanglement across different partitions.&lt;/p&gt;

&lt;p&gt;This picture naturally explains why transport-dominated systems like MBL exhibit their largest observable entanglement growth at intermediate initial entanglement.&lt;/p&gt;




&lt;h2&gt;
  
  
  A new perspective on entanglement dynamics
&lt;/h2&gt;

&lt;p&gt;The broader message is that entanglement dynamics is not just about &lt;strong&gt;creating&lt;/strong&gt; quantum information.&lt;/p&gt;

&lt;p&gt;It is equally about &lt;strong&gt;processing, redistributing, and transporting&lt;/strong&gt; the quantum information that already exists.&lt;/p&gt;

&lt;p&gt;The traditional "start from product states" paradigm has taught us a great deal, but it captures only one half of the story.&lt;/p&gt;

&lt;p&gt;Distinguishing between &lt;strong&gt;entanglement generation&lt;/strong&gt; and &lt;strong&gt;entanglement transport&lt;/strong&gt; provides a more complete framework for understanding quantum dynamics across thermalizing systems, free fermions, many-body localization, and quantum circuits.&lt;/p&gt;

&lt;p&gt;Beyond offering a conceptual picture, this framework makes concrete predictions that can be tested on today's quantum simulation platforms, providing new experimental probes of nonequilibrium quantum dynamics.&lt;/p&gt;




&lt;h2&gt;
  
  
  Reference
&lt;/h2&gt;

&lt;p&gt;&lt;strong&gt;Entanglement growth from entangled states: A unified perspective on entanglement generation and transport&lt;/strong&gt;&lt;/p&gt;

&lt;p&gt;Chun-Yue Zhang, Zi-Xiang Li, and Shi-Xin Zhang&lt;/p&gt;

&lt;p&gt;&lt;em&gt;Physical Review Letters&lt;/em&gt; &lt;strong&gt;137&lt;/strong&gt;, 020404 (2026)&lt;/p&gt;

</description>
      <category>quantum</category>
    </item>
    <item>
      <title>Can AI Really Write Quantum Computing Code? Introducing ORBIT-Q: A Dual-Axis Benchmark for AI Agents and Quantum Software Frameworks</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Wed, 08 Jul 2026 05:35:39 +0000</pubDate>
      <link>https://dev.to/refractionray/can-ai-really-write-quantum-computing-code-introducing-orbit-q-a-dual-axis-benchmark-for-ai-3276</link>
      <guid>https://dev.to/refractionray/can-ai-really-write-quantum-computing-code-introducing-orbit-q-a-dual-axis-benchmark-for-ai-3276</guid>
      <description>&lt;p&gt;Large language models have become remarkably good at writing conventional software. Frameworks like Codex, Claude Code, and other coding agents can already solve a large fraction of real-world programming tasks.&lt;/p&gt;

&lt;p&gt;But scientific programming—especially quantum computing—is a very different challenge.&lt;/p&gt;

&lt;p&gt;A quantum program is not simply expected to compile and produce the correct output. It must also preserve physical correctness, maintain differentiability, respect algorithmic constraints, and often achieve high computational performance. Traditional software benchmarks rarely capture these requirements.&lt;/p&gt;

&lt;p&gt;To evaluate how AI agents perform in this setting, we developed &lt;strong&gt;ORBIT-Q (Open Research Benchmark for Integrated Tasks in Quantum Computing)&lt;/strong&gt;, a benchmark specifically designed for autonomous scientific programming in quantum computing.&lt;/p&gt;




&lt;h2&gt;
  
  
  Why Existing Benchmarks Are Not Enough
&lt;/h2&gt;

&lt;p&gt;In conventional coding benchmarks, passing unit tests is often sufficient.&lt;/p&gt;

&lt;p&gt;In scientific computing, however, an implementation may pass numerical tests while still being fundamentally wrong.&lt;/p&gt;

&lt;p&gt;During our experiments we frequently observed behaviors such as:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;
&lt;strong&gt;Framework bypassing.&lt;/strong&gt; Instead of using the requested quantum framework, the agent secretly reconstructs the computation with NumPy or JAX tensor operations.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;Broken differentiability.&lt;/strong&gt; The generated code produces correct numbers but destroys the end-to-end automatic differentiation pipeline.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;Violation of physical assumptions.&lt;/strong&gt; The implementation changes the intended mathematical or physical problem while still appearing to "work."&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;These failures are difficult to detect using standard execution-based evaluation alone.&lt;/p&gt;

&lt;p&gt;Scientific programming therefore requires evaluation beyond correctness—it requires semantic verification.&lt;/p&gt;




&lt;h2&gt;
  
  
  ORBIT-Q: A Dual-Axis Benchmark
&lt;/h2&gt;

&lt;p&gt;ORBIT-Q consists of &lt;strong&gt;12 challenging research-level quantum programming tasks&lt;/strong&gt;, covering representative workloads in quantum simulation, quantum machine learning, tensor network algorithms, optimization, and automatic differentiation.&lt;/p&gt;

&lt;p&gt;Its key idea is a &lt;strong&gt;dual-axis evaluation protocol&lt;/strong&gt;.&lt;/p&gt;

&lt;h3&gt;
  
  
  Axis 1: Agent Evaluation
&lt;/h3&gt;

&lt;p&gt;Keep the quantum framework fixed and compare different AI agents.&lt;/p&gt;

&lt;p&gt;This measures how well various models (GPT, Claude, etc.) can solve scientific programming tasks under identical software environments.&lt;/p&gt;

&lt;h3&gt;
  
  
  Axis 2: Framework Evaluation
&lt;/h3&gt;

&lt;p&gt;Keep the AI agent fixed and compare different quantum software frameworks.&lt;/p&gt;

&lt;p&gt;This evaluates not only functionality and runtime performance, but also something increasingly important in the AI era:&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;&lt;strong&gt;How AI-friendly is a software framework?&lt;/strong&gt;&lt;/p&gt;
&lt;/blockquote&gt;

&lt;p&gt;A framework with discoverable APIs, composable abstractions, and consistent interfaces allows autonomous agents to generate substantially better solutions.&lt;/p&gt;

&lt;p&gt;&lt;a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2F98ci0c4a969j2v88cvl0.webp" class="article-body-image-wrapper"&gt;&lt;img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2F98ci0c4a969j2v88cvl0.webp" alt="Dual axis eval" width="799" height="267"&gt;&lt;/a&gt;&lt;/p&gt;




&lt;h2&gt;
  
  
  Preventing "Cheating"
&lt;/h2&gt;

&lt;p&gt;To ensure generated solutions genuinely use the intended framework, ORBIT-Q employs a three-stage verification pipeline:&lt;/p&gt;

&lt;ol&gt;
&lt;li&gt;&lt;strong&gt;Deterministic functional testing&lt;/strong&gt;&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;LLM-based source-level semantic auditing&lt;/strong&gt;, designed to detect framework bypassing and other implementation shortcuts&lt;/li&gt;
&lt;li&gt;&lt;strong&gt;Expert manual review&lt;/strong&gt;&lt;/li&gt;
&lt;/ol&gt;

&lt;p&gt;This combination substantially reduces false positives that commonly appear in conventional coding benchmarks.&lt;/p&gt;




&lt;h2&gt;
  
  
  Results: Which Frameworks Work Best?
&lt;/h2&gt;

&lt;p&gt;Using the same coding agent (Codex + GPT-5.5), we evaluated several mainstream quantum software frameworks.&lt;/p&gt;

&lt;p&gt;TensorCircuit-NG achieved the strongest overall performance, successfully solving &lt;strong&gt;10 out of 12&lt;/strong&gt; benchmark tasks while also delivering significantly faster execution than competing frameworks.&lt;/p&gt;

&lt;p&gt;For comparison:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;TensorCircuit-NG: &lt;strong&gt;10 / 12&lt;/strong&gt;
&lt;/li&gt;
&lt;li&gt;PennyLane: &lt;strong&gt;8 / 12&lt;/strong&gt;
&lt;/li&gt;
&lt;li&gt;TorchQuantum: &lt;strong&gt;4 / 12&lt;/strong&gt;
&lt;/li&gt;
&lt;li&gt;MindQuantum: &lt;strong&gt;4 / 12&lt;/strong&gt;
&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;Beyond success rate, TensorCircuit-NG consistently produced solutions that executed several times faster than those generated for other frameworks.&lt;/p&gt;

&lt;p&gt;This suggests that framework design has a substantial impact on autonomous scientific programming—not only for human developers but also for AI agents.&lt;/p&gt;

&lt;p&gt;&lt;a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fjbqylleq3g2ezzaozrrp.webp" class="article-body-image-wrapper"&gt;&lt;img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fjbqylleq3g2ezzaozrrp.webp" alt="Solution succeed rate and efficiency" width="800" height="370"&gt;&lt;/a&gt;&lt;/p&gt;




&lt;h2&gt;
  
  
  Results: Which AI Agents Perform Best?
&lt;/h2&gt;

&lt;p&gt;Under the TensorCircuit-NG environment, the leading coding agents achieved:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;Codex + GPT-5.5: &lt;strong&gt;10 / 12&lt;/strong&gt;
&lt;/li&gt;
&lt;li&gt;Claude Code + Opus 4.8: &lt;strong&gt;9 / 12&lt;/strong&gt;
&lt;/li&gt;
&lt;li&gt;Claude Code + Sonnet 4.6: &lt;strong&gt;7 / 12&lt;/strong&gt;
&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;Although these results are encouraging, a significant gap remains between AI-generated solutions and expert-written implementations.&lt;/p&gt;

&lt;p&gt;Human experts solved &lt;strong&gt;all 12 tasks&lt;/strong&gt;, while producing implementations that were typically &lt;strong&gt;more than twice as efficient&lt;/strong&gt; as the strongest AI-generated solutions.&lt;/p&gt;

&lt;p&gt;Current AI systems are becoming effective research assistants, but they are still far from replacing domain experts in scientific software development.&lt;/p&gt;




&lt;h2&gt;
  
  
  An Unexpected Observation: Safety False Positives
&lt;/h2&gt;

&lt;p&gt;One particularly interesting finding was unrelated to quantum computing itself.&lt;/p&gt;

&lt;p&gt;During evaluation with Claude Code + Opus 4.8, two benchmark tasks failed—not because the model lacked the necessary programming capability, but because the interaction was interrupted by &lt;strong&gt;Cybersecurity Refusals&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;These tasks involved no networking, no external services, and no security-sensitive objectives. They consisted solely of local quantum programming and framework exploration.&lt;/p&gt;

&lt;p&gt;This illustrates an often-overlooked issue in agent evaluation:&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;Product-level safety policies can significantly affect end-to-end task reliability, even when the underlying model is technically capable of solving the problem.&lt;/p&gt;
&lt;/blockquote&gt;

&lt;p&gt;For autonomous scientific workflows, reliability depends not only on model intelligence but also on surrounding product behavior.&lt;/p&gt;




&lt;h2&gt;
  
  
  The Economics of Scientific AI
&lt;/h2&gt;

&lt;p&gt;Another interesting lesson concerns inference cost.&lt;/p&gt;

&lt;p&gt;Developers often compare models by token price alone.&lt;/p&gt;

&lt;p&gt;Our experiments suggest this can be misleading.&lt;/p&gt;

&lt;p&gt;Lower-cost models frequently require many more iterations because they generate incorrect implementations, encounter execution failures, or repeatedly need debugging.&lt;/p&gt;

&lt;p&gt;Consequently, obtaining one &lt;strong&gt;successful scientific solution&lt;/strong&gt; may consume substantially more time and tokens than using a stronger (but more expensive) model.&lt;/p&gt;

&lt;p&gt;For scientific programming, a more meaningful metric may be:&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;&lt;strong&gt;Cost per successful scientific solution&lt;/strong&gt;&lt;/p&gt;
&lt;/blockquote&gt;

&lt;p&gt;rather than simply cost per token.&lt;/p&gt;

&lt;p&gt;&lt;a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fnyku8am2z750tuev7r3i.webp" class="article-body-image-wrapper"&gt;&lt;img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fnyku8am2z750tuev7r3i.webp" alt="Agent efficiency and cost" width="800" height="564"&gt;&lt;/a&gt;&lt;/p&gt;




&lt;h2&gt;
  
  
  Looking Ahead
&lt;/h2&gt;

&lt;p&gt;As AI agents become increasingly integrated into scientific research, software frameworks will need to evolve as well.&lt;/p&gt;

&lt;p&gt;Future scientific software should not only be easy for researchers to use—it should also be easy for autonomous agents to understand, compose, and optimize.&lt;/p&gt;

&lt;p&gt;Although ORBIT-Q focuses on quantum computing, we believe its evaluation methodology can be generalized to many areas of scientific computing where correctness, semantics, differentiability, and performance all matter.&lt;/p&gt;

&lt;p&gt;If AI is going to become a true collaborator in scientific discovery, we need benchmarks that measure much more than whether code simply runs.&lt;/p&gt;




&lt;p&gt;&lt;strong&gt;Paper&lt;/strong&gt;&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;ORBIT-Q: Dual-axis Benchmarking of Autonomous Agents in Scientific Quantum Programming&lt;/strong&gt;&lt;/p&gt;

&lt;p&gt;Shi-Xin Zhang and Yu-Qin Chen&lt;/p&gt;

&lt;p&gt;arXiv:2607.03105&lt;/p&gt;

&lt;p&gt;The benchmark, evaluation framework, and source code are all open source at GitHub: &lt;a href="https://github.com/sxzgroup/ORBIT-Q" rel="noopener noreferrer"&gt;https://github.com/sxzgroup/ORBIT-Q&lt;/a&gt; and the accompanying webpage: &lt;a href="https://sxzgroup.github.io/ORBIT-Q/" rel="noopener noreferrer"&gt;https://sxzgroup.github.io/ORBIT-Q/&lt;/a&gt;.&lt;/p&gt;

</description>
      <category>ai</category>
      <category>agents</category>
      <category>quantum</category>
    </item>
    <item>
      <title>The "Secret of Staying Young" in Quantum Neural Networks</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Tue, 07 Jul 2026 06:57:41 +0000</pubDate>
      <link>https://dev.to/refractionray/the-secret-of-staying-young-in-quantum-neural-networks-knl</link>
      <guid>https://dev.to/refractionray/the-secret-of-staying-young-in-quantum-neural-networks-knl</guid>
      <description>&lt;p&gt;&lt;em&gt;How quantum geometry helps preserve learning ability in continual learning&lt;/em&gt;&lt;/p&gt;

&lt;p&gt;In quantum machine learning, new models are often evaluated by how much they improve benchmark accuracy over classical baselines. These quantitative gains, however, are frequently fragile. They can depend heavily on the choice of baseline models, hyperparameters, or other experimental details.&lt;/p&gt;

&lt;p&gt;A more fundamental question is whether &lt;strong&gt;quantum and classical learning systems exhibit qualitatively different learning dynamics&lt;/strong&gt;. Such structural differences, if they exist, reveal something deeper than a few percentage points of accuracy—they provide insight into the underlying mechanisms of learning itself.&lt;/p&gt;

&lt;p&gt;A recent breakthrough study published in &lt;strong&gt;PRX Quantum by Yu-Qin Chen of the Graduate School of the Chinese Academy of Sciences and Shi-Xin Zhang of the Institute of Physics, Chinese Academy of Sciences&lt;/strong&gt;, explores this question from the perspective of &lt;strong&gt;continual learning&lt;/strong&gt;. Instead of asking whether quantum neural networks achieve higher accuracy, the work asks:&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;&lt;strong&gt;Can quantum neural networks preserve their ability to learn over long periods of continual training? If so, why?&lt;/strong&gt;&lt;/p&gt;
&lt;/blockquote&gt;

&lt;p&gt;The answer turns out to reveal a surprising geometric advantage rooted in quantum mechanics itself.&lt;/p&gt;




&lt;h2&gt;
  
  
  AI's Midlife Crisis: Losing the Ability to Learn
&lt;/h2&gt;

&lt;p&gt;Continual learning aims to build models that, much like humans, continuously accumulate knowledge while adapting to new tasks and changing environments.&lt;/p&gt;

&lt;p&gt;Historically, research has focused on &lt;strong&gt;catastrophic forgetting&lt;/strong&gt;—the tendency of neural networks to overwrite previously learned knowledge when learning new tasks.&lt;/p&gt;

&lt;p&gt;In recent years, however, researchers have recognized another equally important challenge.&lt;/p&gt;

&lt;p&gt;As training continues across many tasks, models gradually become &lt;strong&gt;less capable of learning new information&lt;/strong&gt;. Their parameters become increasingly difficult to update, gradients become less informative, and adaptation slows dramatically.&lt;/p&gt;

&lt;p&gt;This phenomenon is known as &lt;strong&gt;loss of plasticity&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;Among the earliest researchers to emphasize its importance was reinforcement learning pioneer &lt;strong&gt;Richard Sutton&lt;/strong&gt;, who argued that for long-running learning systems, catastrophic forgetting and loss of plasticity are two complementary problems:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;Catastrophic forgetting determines how well a model retains old knowledge.&lt;/li&gt;
&lt;li&gt;Loss of plasticity determines how well it can acquire new knowledge.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;An intuitive analogy is that the model gradually "ages." Although it accumulates more experience, it becomes increasingly resistant to learning anything new.&lt;/p&gt;




&lt;p&gt;&lt;a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fsf5grcsigudsi0xiu1r8.png" class="article-body-image-wrapper"&gt;&lt;img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fsf5grcsigudsi0xiu1r8.png" alt="QNN keeps learning ability across different tasks" width="800" height="458"&gt;&lt;/a&gt;&lt;/p&gt;

&lt;h2&gt;
  
  
  Do Quantum Models Age More Slowly?
&lt;/h2&gt;

&lt;p&gt;The natural question is whether quantum neural networks suffer from the same phenomenon.&lt;/p&gt;

&lt;p&gt;The study first addressed a simple question:&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;&lt;strong&gt;Do quantum neural networks preserve plasticity better than classical neural networks?&lt;/strong&gt;&lt;/p&gt;
&lt;/blockquote&gt;

&lt;p&gt;The answer appears to be &lt;strong&gt;yes&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;Across continual learning experiments involving &lt;strong&gt;more than 3,000 sequential tasks&lt;/strong&gt;, a remarkably consistent pattern emerged.&lt;/p&gt;

&lt;p&gt;Classical neural networks steadily lost their learning ability as training progressed.&lt;/p&gt;

&lt;p&gt;Quantum neural networks, in contrast, maintained a much higher level of plasticity throughout long training sequences.&lt;/p&gt;

&lt;p&gt;But observing the phenomenon is only the beginning.&lt;/p&gt;

&lt;p&gt;The more interesting question is:&lt;/p&gt;

&lt;blockquote&gt;
&lt;p&gt;&lt;strong&gt;Why does this happen?&lt;/strong&gt;&lt;/p&gt;
&lt;/blockquote&gt;




&lt;h2&gt;
  
  
  Geometry Matters
&lt;/h2&gt;

&lt;p&gt;To understand the difference, consider where the parameters of a neural network live.&lt;/p&gt;

&lt;p&gt;Classical neural network weights inhabit ordinary &lt;strong&gt;Euclidean space&lt;/strong&gt;. In principle, parameter norms can grow without bound.&lt;/p&gt;

&lt;p&gt;During prolonged continual training, optimization often drives these parameters toward increasingly large magnitudes.&lt;/p&gt;

&lt;p&gt;Initially, this helps fit the data.&lt;/p&gt;

&lt;p&gt;Eventually, however, several undesirable effects emerge:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;neurons become increasingly saturated,&lt;/li&gt;
&lt;li&gt;effective gradients shrink,&lt;/li&gt;
&lt;li&gt;parameter updates become harder,&lt;/li&gt;
&lt;li&gt;the trace of the Fisher Information Matrix steadily decreases.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;Together, these effects gradually reduce the model's ability to adapt to new tasks.&lt;/p&gt;

&lt;p&gt;This suggests that &lt;strong&gt;loss of plasticity is fundamentally connected to the geometry of the parameter space.&lt;/strong&gt;&lt;/p&gt;




&lt;h2&gt;
  
  
  Why Quantum Neural Networks Behave Differently
&lt;/h2&gt;

&lt;p&gt;Quantum neural networks follow an entirely different geometric trajectory.&lt;/p&gt;

&lt;p&gt;The reason is not a specially designed continual-learning algorithm.&lt;/p&gt;

&lt;p&gt;Instead, it originates from one of the most fundamental principles of quantum mechanics.&lt;/p&gt;

&lt;p&gt;Quantum evolution is described by &lt;strong&gt;unitary transformations&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;Mathematically, the trainable parameters correspond to rotations on &lt;strong&gt;compact Lie groups&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;Unlike Euclidean space, these parameter manifolds are &lt;strong&gt;compact&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;Parameters can continue evolving indefinitely, but they cannot drift arbitrarily far away.&lt;/p&gt;

&lt;p&gt;This geometric constraint naturally prevents the unbounded parameter growth observed in classical networks.&lt;/p&gt;

&lt;p&gt;As a result,&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;gradients remain in a healthy range,&lt;/li&gt;
&lt;li&gt;parameter norms stay bounded,&lt;/li&gt;
&lt;li&gt;the Fisher Information Matrix remains active,&lt;/li&gt;
&lt;li&gt;and the network continues to retain the ability to learn new tasks.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;In other words, the advantage of quantum models may not come solely from having richer computational representations.&lt;/p&gt;

&lt;p&gt;It may also arise from the &lt;strong&gt;geometry imposed by the laws of quantum physics&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;Rather than expanding without limit, quantum parameters evolve on a compact manifold whose structure naturally protects learning plasticity over time.&lt;/p&gt;

&lt;p&gt;This geometric explanation is arguably more interesting than reporting another benchmark improvement.&lt;/p&gt;

&lt;p&gt;Instead of asking whether one model wins by a few percentage points on a particular dataset, it asks whether &lt;strong&gt;quantum and classical learning systems obey fundamentally different learning dynamics&lt;/strong&gt; during long-term adaptation.&lt;/p&gt;




&lt;h2&gt;
  
  
  From Theory to Large-Scale Validation
&lt;/h2&gt;

&lt;p&gt;A theoretical explanation is only convincing if it survives large-scale empirical testing.&lt;/p&gt;

&lt;p&gt;To validate the proposed mechanism, the authors constructed multiple continual learning benchmarks involving&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;more than &lt;strong&gt;3,000 sequential learning tasks&lt;/strong&gt;,&lt;/li&gt;
&lt;li&gt;quantum circuits with depths up to &lt;strong&gt;30 layers&lt;/strong&gt;,&lt;/li&gt;
&lt;li&gt;and over &lt;strong&gt;4,000 trainable quantum parameters&lt;/strong&gt;.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;These experiments are considerably more demanding than conventional machine learning benchmarks.&lt;/p&gt;

&lt;p&gt;Each configuration effectively requires training thousands of quantum neural networks while continuously monitoring internal quantities such as gradient statistics and the Fisher Information Matrix throughout optimization.&lt;/p&gt;

&lt;p&gt;Such experiments would be prohibitively slow—or simply infeasible—using many conventional quantum software frameworks.&lt;/p&gt;

&lt;p&gt;The computational foundation of this work therefore relied heavily on &lt;strong&gt;TensorCircuit-NG&lt;/strong&gt;, an open-source quantum computing framework that combines tensor-network simulation, automatic differentiation, and high-performance GPU acceleration. These capabilities make long-horizon, large-scale continual learning experiments computationally practical.&lt;/p&gt;

&lt;p&gt;&lt;a href="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fhkseifvliw8v4vtxwdha.png" class="article-body-image-wrapper"&gt;&lt;img src="https://media2.dev.to/dynamic/image/width=800%2Cheight=%2Cfit=scale-down%2Cgravity=auto%2Cformat=auto/https%3A%2F%2Fdev-to-uploads.s3.us-east-2.amazonaws.com%2Fuploads%2Farticles%2Fhkseifvliw8v4vtxwdha.png" alt="QNN is also superior in RL settings" width="800" height="763"&gt;&lt;/a&gt;&lt;/p&gt;




&lt;h2&gt;
  
  
  A Different Perspective on Quantum Advantage
&lt;/h2&gt;

&lt;p&gt;This work does &lt;strong&gt;not&lt;/strong&gt; claim that quantum neural networks have solved continual learning.&lt;/p&gt;

&lt;p&gt;Catastrophic forgetting still exists, and many questions about memory retention, stability, and continual adaptation remain open.&lt;/p&gt;

&lt;p&gt;Instead, the paper offers a different perspective on &lt;strong&gt;quantum advantage&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;Discussions of quantum machine learning often emphasize computational speedups or asymptotic complexity advantages.&lt;/p&gt;

&lt;p&gt;But real intelligent systems require more than fast learning.&lt;/p&gt;

&lt;p&gt;They must also &lt;strong&gt;continue learning over time&lt;/strong&gt;.&lt;/p&gt;

&lt;p&gt;Continual learning requires both remembering what has already been learned and remaining capable of acquiring new knowledge.&lt;/p&gt;

&lt;p&gt;Catastrophic forgetting addresses the first challenge.&lt;/p&gt;

&lt;p&gt;Loss of plasticity addresses the second.&lt;/p&gt;

&lt;p&gt;Both are essential.&lt;/p&gt;

&lt;p&gt;If quantum neural networks can naturally preserve their capacity to learn throughout long-term adaptation—not because of additional engineering tricks, but because of the geometry dictated by quantum mechanics—then this "ageless" plasticity may represent a compelling and fundamentally different form of quantum advantage.&lt;/p&gt;




&lt;p&gt;&lt;strong&gt;Reference&lt;/strong&gt;&lt;/p&gt;

&lt;p&gt;Chen, Y.-Q., &amp;amp; Zhang, S.-X. (2026). &lt;em&gt;Intrinsic Preservation of Plasticity in Continual Quantum Learning&lt;/em&gt;. &lt;strong&gt;PRX Quantum&lt;/strong&gt;, 7, 033003.&lt;/p&gt;

</description>
      <category>quantum</category>
    </item>
    <item>
      <title>The Two Paradigms of Scientific Computing Agents: Abstraction, Openness, and "The Bitter Lesson"</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Mon, 22 Jun 2026 13:05:07 +0000</pubDate>
      <link>https://dev.to/refractionray/the-two-paradigms-of-scientific-computing-agents-abstraction-openness-and-the-bitter-lesson-8ck</link>
      <guid>https://dev.to/refractionray/the-two-paradigms-of-scientific-computing-agents-abstraction-openness-and-the-bitter-lesson-8ck</guid>
      <description>&lt;p&gt;In recent years, the rapid evolution of Large Language Models (LLMs) has turned "AI + Scientific Computing" into a highly active frontier. Whether in molecular dynamics, material and drug design, or quantum computing, numerous platforms are attempting to bridge natural language interfaces with rigorous scientific computation.&lt;/p&gt;

&lt;p&gt;From a user experience perspective, this approach significantly lowers the barrier to entry, allowing non-experts to breeze through standardized experimental workflows. However, when we shift our focus from &lt;em&gt;"Can it run a standard experiment quickly?"&lt;/em&gt; to &lt;em&gt;"Does it support open-ended scientific exploration?"&lt;/em&gt;, a stark architectural divide emerges regarding abstraction boundaries and system openness.&lt;/p&gt;

&lt;p&gt;Currently, Scientific Computing Agent systems can be broadly categorized into two technical paradigms:&lt;/p&gt;

&lt;ol&gt;
&lt;li&gt;
&lt;strong&gt;Encapsulated Systems:&lt;/strong&gt; Running in controlled cloud sandboxes, these systems typically provide pre-configured, templated workflows accessible via a Web UI.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;Open &amp;amp; Programmable Systems:&lt;/strong&gt; Operating within general-purpose computing environments, these systems (like Claude Code or Codex) integrate deeply with code repositories, runtimes, and external toolchains.&lt;/li&gt;
&lt;/ol&gt;

&lt;p&gt;While both rely on conversational interfaces, their core difference lies in their habitat: is the Agent living in a closed ecosystem of cloud templates, or an open, customizable computing space?&lt;/p&gt;

&lt;h2&gt;
  
  
  Abstraction Boundaries vs. The Space for Innovation
&lt;/h2&gt;

&lt;p&gt;Every software system must strike a balance between ease of use and flexibility. For standardized scientific tasks, encapsulated systems shine. However, when a research question deviates from standard templates, the very abstraction that reduces complexity becomes a bottleneck.&lt;/p&gt;

&lt;p&gt;Here is a clear comparison of the two paradigms:&lt;/p&gt;

&lt;div class="table-wrapper-paragraph"&gt;&lt;table&gt;
&lt;thead&gt;
&lt;tr&gt;
&lt;th&gt;Dimension&lt;/th&gt;
&lt;th&gt;Encapsulated Systems&lt;/th&gt;
&lt;th&gt;Open &amp;amp; Programmable Systems&lt;/th&gt;
&lt;/tr&gt;
&lt;/thead&gt;
&lt;tbody&gt;
&lt;tr&gt;
&lt;td&gt;&lt;strong&gt;Representative Examples&lt;/strong&gt;&lt;/td&gt;
&lt;td&gt;Domain-specific Web-based AI platforms&lt;/td&gt;
&lt;td&gt;General-purpose agents like Claude Code, Codex&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;&lt;strong&gt;Execution Environment&lt;/strong&gt;&lt;/td&gt;
&lt;td&gt;Pre-defined, controlled cloud sandboxes; highly templated&lt;/td&gt;
&lt;td&gt;General compute environments (native OS, containers, local/cloud)&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;&lt;strong&gt;Abstraction Boundary&lt;/strong&gt;&lt;/td&gt;
&lt;td&gt;
&lt;strong&gt;High&lt;/strong&gt; (Hides underlying engineering details)&lt;/td&gt;
&lt;td&gt;
&lt;strong&gt;Low&lt;/strong&gt; (Direct access to file systems, low-level compute libraries, and dependencies)&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;&lt;strong&gt;Ideal Use Cases&lt;/strong&gt;&lt;/td&gt;
&lt;td&gt;Education, running standard algorithms, rapid benchmarking&lt;/td&gt;
&lt;td&gt;Exploratory frontier research, highly customized workflows&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;&lt;strong&gt;Handling Novel Problems&lt;/strong&gt;&lt;/td&gt;
&lt;td&gt;Wait for platform updates, or revert to writing code manually&lt;/td&gt;
&lt;td&gt;Break out of the framework; freely compose modules and custom logic&lt;/td&gt;
&lt;/tr&gt;
&lt;/tbody&gt;
&lt;/table&gt;&lt;/div&gt;

&lt;p&gt;Take &lt;strong&gt;quantum computing&lt;/strong&gt; as an example. For standard Variational Quantum Algorithms (VQAs), encapsulated Web platforms can easily handle the entire pipeline—from quantum circuit construction and parameter optimization to result visualization. By condensing complex engineering details into a few pre-built templates, users can complete experiments with minimal cognitive load.&lt;/p&gt;

&lt;p&gt;But the moment a researcher’s needs veer off the beaten path, this abstraction hits a wall. Suppose a researcher wants to combine a novel data encoding method, a highly customized quantum gate structure, and a non-standard loss function. Because this bespoke architecture doesn't map to existing templates, the encapsulated system's API simply rejects it.&lt;/p&gt;

&lt;p&gt;In contrast, &lt;strong&gt;open programmable systems&lt;/strong&gt; support these novel combinations because they don't pre-define the shape of the problem; they merely provide computing primitives. In these environments, circuit construction, training loops, loss functions, and data pipelines are all exposed as raw code. An Agent (or researcher) can freely import new Python modules, alter the training loop, inject custom gradient estimation methods, or couple a quantum simulator with an external data pipeline. Because the system hasn't hardcoded these steps into indivisible blocks, a problem that breaks an encapsulated system is just another day of writing code for an open system.&lt;/p&gt;

&lt;h2&gt;
  
  
  Project Context and the Information Horizon
&lt;/h2&gt;

&lt;p&gt;If architecture forms the skeleton of a system, context forms the Agent's horizon. The quality of an Agent's reasoning is inextricably tied to the scope of information it can access.&lt;/p&gt;

&lt;p&gt;In real-world scientific computing, the "state" of a project is never just a few chat prompts or isolated data uploads. It is a massive, ongoing web of information: repository directory structures, historical scripts, local test datasets, related PDF papers, version control histories, and past error logs.&lt;/p&gt;

&lt;p&gt;Constrained by cloud sandbox isolation, an encapsulated Agent's horizon is usually limited to the current ephemeral session; its understanding of state is fragmented. Conversely, in an open programmable system, an Agent like Claude Code operates as a first-class citizen within the compute environment. It can directly read the real-time state of the entire project directory. If a user asks to tweak an initialization parameter based on the last run, the open Agent can fetch historical logs, diff code versions, and execute reliable reasoning backed by full project context.&lt;/p&gt;

&lt;p&gt;The difference is fundamental: is the Agent trapped in a single, isolated interaction, or is it embedded in the continuous information network of a real research project?&lt;/p&gt;

&lt;h2&gt;
  
  
  From Code Generation to Workflow Orchestration
&lt;/h2&gt;

&lt;p&gt;Context dictates reasoning, but action drives results. Once an Agent grasps the global state, its capabilities expand from mere code generation to system-level orchestration. This is the second great divide between the two paradigms: the breadth of agency.&lt;/p&gt;

&lt;p&gt;Scientific computing rarely stops at writing a single algorithm script. It is usually a complex workflow spanning multiple independent tools. Under the open paradigm, an Agent doesn't just write logic using native frameworks; it executes system-level operations. It can SSH into High-Performance Computing (HPC) clusters to submit jobs, set up cron jobs to monitor GPU utilization, read stack traces to auto-retry crashed experiments, clean data post-run, generate charts, and even draft the initial manuscript.&lt;/p&gt;

&lt;p&gt;While encapsulated systems confine the Agent to a proprietary loop, open systems grant Agents cross-platform, cross-tool autonomy, turning them into genuine collaborators.&lt;/p&gt;

&lt;h2&gt;
  
  
  General Beats Specialized: The "Bitter Lesson" in the Agent Era
&lt;/h2&gt;

&lt;p&gt;This transition from encapsulated to open, and from specialized cloud platforms to general programming environments, perfectly echoes Richard Sutton's famous essay, &lt;em&gt;"The Bitter Lesson"&lt;/em&gt;. Sutton observed that throughout AI history, specialized methods meticulously hand-crafted using human domain knowledge are ultimately surpassed by general methods that leverage massive computation.&lt;/p&gt;

&lt;p&gt;This philosophy holds entirely true for AI in scientific computing. Today, many platforms pour immense resources into building highly specialized Agents for niche domains, designing bespoke UIs and rigid workflow templates. In the short term, these make for incredibly smooth demos.&lt;/p&gt;

&lt;p&gt;However, as the capabilities of foundation models scale exponentially, general-purpose Agents (like Claude Code or Codex) are becoming overwhelmingly powerful. They don't need a specialized UI wrapper. Drop them into a standard, open computing environment, and their generalized reasoning allows them to read domain documentation on the fly, call low-level scientific libraries, and independently orchestrate complex analysis.&lt;/p&gt;

&lt;p&gt;The meticulously hardcoded workflows of domain-specific Agents risk rapid obsolescence. Often, their domain capabilities struggle to keep pace with the generalized leaps made by foundation models.&lt;/p&gt;

&lt;p&gt;Recognizing this, a new generation of scientific computing frameworks is aligning with the open paradigm. For instance, in the quantum computing space, &lt;strong&gt;TensorCircuit-NG&lt;/strong&gt; is a prime example of embracing the general Agent route. It abandons traditional closed-platform UI wrappers in favor of a native AI stack, offering hyper-performance low-level primitives alongside out-of-the-box skill suites. Its core design philosophy is simple: empower general-purpose Agents to freely explore and orchestrate complex science within an open environment.&lt;/p&gt;

&lt;h2&gt;
  
  
  Conclusion
&lt;/h2&gt;

&lt;p&gt;Encapsulated and open programmable systems represent two distinct technological philosophies. The former lowers the barrier to entry via high-level abstraction, playing a crucial role in education and standard reproducibility.&lt;/p&gt;

&lt;p&gt;However, in uncharted, fast-moving frontier sciences, maintaining system openness and generality is paramount. Allowing increasingly capable general Agents to dive deep into real, low-level engineering environments—breaking free from rigid abstraction boundaries—is the most sustainable path for AI to truly revolutionize scientific computing.&lt;/p&gt;

</description>
      <category>ai</category>
      <category>agents</category>
      <category>hpc</category>
    </item>
    <item>
      <title>PyTrees Are Not One Thing: JAX, PyTorch, and TensorFlow Compared</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Fri, 12 Jun 2026 06:18:20 +0000</pubDate>
      <link>https://dev.to/refractionray/pytrees-are-not-one-thing-jax-pytorch-and-tensorflow-compared-hjh</link>
      <guid>https://dev.to/refractionray/pytrees-are-not-one-thing-jax-pytorch-and-tensorflow-compared-hjh</guid>
      <description>&lt;p&gt;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.&lt;/p&gt;

&lt;p&gt;This note compares JAX &lt;code&gt;jax.tree_util&lt;/code&gt;, PyTorch &lt;code&gt;torch.utils._pytree&lt;/code&gt;, and TensorFlow &lt;code&gt;tf.nest&lt;/code&gt;. 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: &lt;code&gt;_pytree.tree_map&lt;/code&gt; in 2.2.2 accepts only one pytree, while 2.12.0 supports multiple pytrees and behaves much closer to JAX prefix-style mapping.&lt;/p&gt;

&lt;p&gt;The short version: JAX treats pytrees as a transformation language, PyTorch is converging toward that model in &lt;code&gt;torch.func&lt;/code&gt;, and TensorFlow exposes a broader nested-structure utility through &lt;code&gt;tf.nest&lt;/code&gt;. Those differences show up exactly where backend-agnostic libraries usually hurt: &lt;code&gt;None&lt;/code&gt;, dictionary order, custom containers, &lt;code&gt;tree_map&lt;/code&gt;, autodiff, and vectorization.&lt;/p&gt;

&lt;h2&gt;
  
  
  The Shape Of The APIs
&lt;/h2&gt;

&lt;p&gt;The three APIs have the same surface story but not the same contract.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="kn"&gt;from&lt;/span&gt; &lt;span class="n"&gt;jax&lt;/span&gt; &lt;span class="kn"&gt;import&lt;/span&gt; &lt;span class="n"&gt;tree_util&lt;/span&gt; &lt;span class="k"&gt;as&lt;/span&gt; &lt;span class="n"&gt;jtu&lt;/span&gt;
&lt;span class="kn"&gt;from&lt;/span&gt; &lt;span class="n"&gt;torch.utils&lt;/span&gt; &lt;span class="kn"&gt;import&lt;/span&gt; &lt;span class="n"&gt;_pytree&lt;/span&gt; &lt;span class="k"&gt;as&lt;/span&gt; &lt;span class="n"&gt;tpu&lt;/span&gt;
&lt;span class="kn"&gt;import&lt;/span&gt; &lt;span class="n"&gt;tensorflow&lt;/span&gt; &lt;span class="k"&gt;as&lt;/span&gt; &lt;span class="n"&gt;tf&lt;/span&gt;

&lt;span class="n"&gt;leaves&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;treedef&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_unflatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;treedef&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;leaves&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;f&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt;&lt;span class="n"&gt;trees&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;

&lt;span class="n"&gt;leaves&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;spec&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_unflatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;leaves&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;spec&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;f&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;          &lt;span class="c1"&gt;# PyTorch 2.2.2
&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;f&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt;&lt;span class="n"&gt;trees&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;        &lt;span class="c1"&gt;# PyTorch 2.12.0
&lt;/span&gt;
&lt;span class="n"&gt;leaves&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;pack_sequence_as&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;structure&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;leaves&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;map_structure&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;f&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="o"&gt;*&lt;/span&gt;&lt;span class="n"&gt;structures&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Flattening means "which objects are leaves?" Unflattening means "what metadata is needed to reconstruct the original container?" Mapping means "what does it mean for several structures to match?" Those three questions are where the frameworks diverge.&lt;/p&gt;

&lt;p&gt;JAX calls its structure descriptor a &lt;code&gt;PyTreeDef&lt;/code&gt;, so &lt;code&gt;treedef&lt;/code&gt; is the conventional variable name. PyTorch calls the analogous descriptor a &lt;code&gt;TreeSpec&lt;/code&gt;, so examples and internals often name it &lt;code&gt;spec&lt;/code&gt;. Conceptually they play the same role: they describe the container skeleton and the metadata needed to rebuild it from a flat leaf list. TensorFlow's &lt;code&gt;tf.nest&lt;/code&gt; does not return a separate treedef object from &lt;code&gt;flatten&lt;/code&gt;; instead, &lt;code&gt;pack_sequence_as&lt;/code&gt; takes an existing nested &lt;code&gt;structure&lt;/code&gt; as the template.&lt;/p&gt;

&lt;p&gt;There is also a small argument-order trap. JAX unflattens as &lt;code&gt;tree_unflatten(treedef, leaves)&lt;/code&gt;, while PyTorch unflattens as &lt;code&gt;tree_unflatten(leaves, spec)&lt;/code&gt;. TensorFlow's equivalent is &lt;code&gt;pack_sequence_as(structure, leaves)&lt;/code&gt;. &lt;/p&gt;

&lt;h2&gt;
  
  
  A Compact Map Of The Differences
&lt;/h2&gt;

&lt;div class="table-wrapper-paragraph"&gt;&lt;table&gt;
&lt;thead&gt;
&lt;tr&gt;
&lt;th&gt;Case&lt;/th&gt;
&lt;th&gt;JAX&lt;/th&gt;
&lt;th&gt;PyTorch &lt;code&gt;_pytree&lt;/code&gt;
&lt;/th&gt;
&lt;th&gt;TensorFlow &lt;code&gt;tf.nest&lt;/code&gt;
&lt;/th&gt;
&lt;/tr&gt;
&lt;/thead&gt;
&lt;tbody&gt;
&lt;tr&gt;
&lt;td&gt;Scalar&lt;/td&gt;
&lt;td&gt;Leaf&lt;/td&gt;
&lt;td&gt;Leaf&lt;/td&gt;
&lt;td&gt;Leaf&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;&lt;code&gt;None&lt;/code&gt;&lt;/td&gt;
&lt;td&gt;Empty pytree, 0 leaves&lt;/td&gt;
&lt;td&gt;Leaf&lt;/td&gt;
&lt;td&gt;Leaf&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;
&lt;code&gt;list&lt;/code&gt;, &lt;code&gt;tuple&lt;/code&gt;
&lt;/td&gt;
&lt;td&gt;Containers&lt;/td&gt;
&lt;td&gt;Containers&lt;/td&gt;
&lt;td&gt;Containers&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;namedtuple&lt;/td&gt;
&lt;td&gt;Container, type-strict&lt;/td&gt;
&lt;td&gt;Container, type-strict&lt;/td&gt;
&lt;td&gt;Container, type-strict&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;plain &lt;code&gt;dict&lt;/code&gt; order&lt;/td&gt;
&lt;td&gt;Sorted keys&lt;/td&gt;
&lt;td&gt;Insertion order&lt;/td&gt;
&lt;td&gt;Sorted-key leaf order&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;
&lt;code&gt;OrderedDict&lt;/code&gt; order&lt;/td&gt;
&lt;td&gt;Insertion order&lt;/td&gt;
&lt;td&gt;Insertion order&lt;/td&gt;
&lt;td&gt;Sorted-key leaf order&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;
&lt;code&gt;defaultdict&lt;/code&gt; order&lt;/td&gt;
&lt;td&gt;Sorted keys&lt;/td&gt;
&lt;td&gt;Insertion order&lt;/td&gt;
&lt;td&gt;Sorted-key leaf order&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;&lt;code&gt;defaultdict.default_factory&lt;/code&gt;&lt;/td&gt;
&lt;td&gt;Preserved&lt;/td&gt;
&lt;td&gt;Preserved&lt;/td&gt;
&lt;td&gt;Preserved&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;custom &lt;code&gt;dict&lt;/code&gt; subclass&lt;/td&gt;
&lt;td&gt;Leaf unless registered&lt;/td&gt;
&lt;td&gt;Leaf unless registered&lt;/td&gt;
&lt;td&gt;Container&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;custom &lt;code&gt;list&lt;/code&gt;/&lt;code&gt;tuple&lt;/code&gt; subclass&lt;/td&gt;
&lt;td&gt;Leaf unless registered&lt;/td&gt;
&lt;td&gt;Leaf unless registered&lt;/td&gt;
&lt;td&gt;Container&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;dataclass instance&lt;/td&gt;
&lt;td&gt;Leaf unless registered&lt;/td&gt;
&lt;td&gt;Leaf unless registered&lt;/td&gt;
&lt;td&gt;Leaf by default&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;multi-arg &lt;code&gt;tree_map&lt;/code&gt;
&lt;/td&gt;
&lt;td&gt;Supported, prefix semantics&lt;/td&gt;
&lt;td&gt;PyTorch 2.2.2: not supported; PyTorch 2.12.0: supported with prefix semantics&lt;/td&gt;
&lt;td&gt;Supported, strict same structure&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;unflatten arity mismatch&lt;/td&gt;
&lt;td&gt;Raises &lt;code&gt;ValueError&lt;/code&gt;
&lt;/td&gt;
&lt;td&gt;Raises &lt;code&gt;ValueError&lt;/code&gt;
&lt;/td&gt;
&lt;td&gt;Raises &lt;code&gt;ValueError&lt;/code&gt;
&lt;/td&gt;
&lt;/tr&gt;
&lt;/tbody&gt;
&lt;/table&gt;&lt;/div&gt;

&lt;p&gt;The rest of the note explains why these rows matter.&lt;/p&gt;

&lt;h2&gt;
  
  
  &lt;code&gt;None&lt;/code&gt;: A Ghost Node In JAX, A Leaf Elsewhere
&lt;/h2&gt;

&lt;p&gt;The cleanest way to feel the philosophical split is &lt;code&gt;None&lt;/code&gt;. In JAX, &lt;code&gt;None&lt;/code&gt; is not a value to map over. It is a zero-leaf structural marker.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# leaves: []
# treedef: PyTreeDef(None)
&lt;/span&gt;
&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;mapped&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# None
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;In PyTorch and TensorFlow, &lt;code&gt;None&lt;/code&gt; is a leaf.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# [None]
&lt;/span&gt;
&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# [None]
&lt;/span&gt;
&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;mapped&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# ("mapped", None)
&lt;/span&gt;
&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;map_structure&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;mapped&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# ("mapped", None)
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;The nested case makes the difference visible:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;3&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;

&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# [1, 3]
&lt;/span&gt;
&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# [1, None, 3]
&lt;/span&gt;
&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# [1, None, 3]
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;If &lt;code&gt;None&lt;/code&gt; means "optional value absent", JAX treats it structurally. If &lt;code&gt;None&lt;/code&gt; means "a value in my tree", PyTorch and TensorFlow are closer to that intuition.&lt;/p&gt;

&lt;h2&gt;
  
  
  Dictionaries: The Same Keys, Different Time Arrows
&lt;/h2&gt;

&lt;p&gt;Plain &lt;code&gt;dict&lt;/code&gt; is a container everywhere, but the traversal order differs. JAX sorts keys, PyTorch follows insertion order, and TensorFlow assigns leaves by sorted keys while preserving the original mapping order when rebuilding.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;}&lt;/span&gt;

&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# [1, 2]   # a, then b
&lt;/span&gt;
&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# [2, 1]   # b, then a
&lt;/span&gt;
&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# [1, 2]   # a, then b
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Replacing the leaves with &lt;code&gt;[10, 20]&lt;/code&gt; shows the reconstruction contract:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="c1"&gt;# JAX
&lt;/span&gt;&lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;20&lt;/span&gt;&lt;span class="p"&gt;}&lt;/span&gt;

&lt;span class="c1"&gt;# PyTorch
&lt;/span&gt;&lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;20&lt;/span&gt;&lt;span class="p"&gt;}&lt;/span&gt;

&lt;span class="c1"&gt;# TensorFlow
&lt;/span&gt;&lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;20&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;}&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;TensorFlow's result is the surprising one on first read. It maps values according to sorted keys, but prints in the original insertion order. The object order and the leaf assignment order are not the same concept.&lt;/p&gt;

&lt;p&gt;Mixed incomparable key types are another consequence of sorting. JAX and TensorFlow fail on &lt;code&gt;{1: "one", "2": "two"}&lt;/code&gt; because &lt;code&gt;1 &amp;lt; "2"&lt;/code&gt; is not defined. PyTorch does not sort and therefore flattens this case in insertion order.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;({&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;one&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;2&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;two&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;})&lt;/span&gt;
&lt;span class="c1"&gt;# ValueError: Comparator raised exception while sorting pytree dictionary keys.
&lt;/span&gt;
&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;flatten&lt;/span&gt;&lt;span class="p"&gt;({&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;one&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;2&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;two&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;})&lt;/span&gt;
&lt;span class="c1"&gt;# TypeError: '&amp;lt;' not supported between instances of 'str' and 'int'
&lt;/span&gt;
&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;({&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;one&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;2&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;two&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;})[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# ["one", "two"]
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;h2&gt;
  
  
  Ordered Containers Are Not Just Dicts With Better Manners
&lt;/h2&gt;

&lt;p&gt;&lt;code&gt;OrderedDict&lt;/code&gt; has explicit order metadata, and JAX treats that metadata as part of the tree structure. PyTorch uses insertion order too. TensorFlow again uses sorted-key leaf assignment.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="kn"&gt;from&lt;/span&gt; &lt;span class="n"&gt;collections&lt;/span&gt; &lt;span class="kn"&gt;import&lt;/span&gt; &lt;span class="n"&gt;OrderedDict&lt;/span&gt;

&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="nc"&gt;OrderedDict&lt;/span&gt;&lt;span class="p"&gt;([(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)])&lt;/span&gt;

&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# [2, 1]
&lt;/span&gt;
&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# [2, 1]
&lt;/span&gt;
&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# [1, 2]
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;All three preserve the &lt;code&gt;OrderedDict&lt;/code&gt; type when rebuilding, but TensorFlow assigns replacement leaves by sorted key:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;pack_sequence_as&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nc"&gt;OrderedDict&lt;/span&gt;&lt;span class="p"&gt;([(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)]),&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;20&lt;/span&gt;&lt;span class="p"&gt;])&lt;/span&gt;
&lt;span class="c1"&gt;# OrderedDict([("b", 20), ("a", 10)])
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Multi-argument mapping reveals another difference. JAX rejects two &lt;code&gt;OrderedDict&lt;/code&gt;s with the same keys but different order because the custom node metadata differs.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;a&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="nc"&gt;OrderedDict&lt;/span&gt;&lt;span class="p"&gt;([(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;)])&lt;/span&gt;
&lt;span class="n"&gt;b&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="nc"&gt;OrderedDict&lt;/span&gt;&lt;span class="p"&gt;([(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;20&lt;/span&gt;&lt;span class="p"&gt;)])&lt;/span&gt;

&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="n"&gt;a&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;b&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# ValueError: Mismatch custom node data: ('b', 'a') != ('a', 'b')
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;TensorFlow accepts this and pairs by key while preserving the first structure's order. PyTorch 2.12.0 also accepts it and returns the same visible result:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="nc"&gt;OrderedDict&lt;/span&gt;&lt;span class="p"&gt;([(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;20&lt;/span&gt;&lt;span class="p"&gt;)),&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;))])&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;h2&gt;
  
  
  &lt;code&gt;defaultdict&lt;/code&gt;: Losing The Type Changes Behavior
&lt;/h2&gt;

&lt;p&gt;&lt;code&gt;defaultdict&lt;/code&gt; is not a decorative subclass. It carries a &lt;code&gt;default_factory&lt;/code&gt;, which changes lookup behavior.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="kn"&gt;from&lt;/span&gt; &lt;span class="n"&gt;collections&lt;/span&gt; &lt;span class="kn"&gt;import&lt;/span&gt; &lt;span class="n"&gt;defaultdict&lt;/span&gt;

&lt;span class="n"&gt;counter&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="nf"&gt;defaultdict&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nb"&gt;int&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;counter&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;missing&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# 0
&lt;/span&gt;
&lt;span class="n"&gt;plain&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;{}&lt;/span&gt;
&lt;span class="n"&gt;plain&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;missing&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# KeyError
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;All three frameworks preserve the &lt;code&gt;default_factory&lt;/code&gt;, but they disagree about leaf order just as with dictionaries.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;tree&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="nf"&gt;defaultdict&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nb"&gt;int&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;})&lt;/span&gt;

&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# [1, 2]
&lt;/span&gt;
&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# [2, 1]
&lt;/span&gt;
&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tree&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# [1, 2]
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Rebuilding with &lt;code&gt;[10, 20]&lt;/code&gt; gives:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="c1"&gt;# JAX
&lt;/span&gt;&lt;span class="nf"&gt;defaultdict&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nb"&gt;int&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;20&lt;/span&gt;&lt;span class="p"&gt;})&lt;/span&gt;

&lt;span class="c1"&gt;# PyTorch
&lt;/span&gt;&lt;span class="nf"&gt;defaultdict&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nb"&gt;int&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;20&lt;/span&gt;&lt;span class="p"&gt;})&lt;/span&gt;

&lt;span class="c1"&gt;# TensorFlow
&lt;/span&gt;&lt;span class="nf"&gt;defaultdict&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nb"&gt;int&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;20&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;})&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;This matters for any pure Python fallback. If it flattens a &lt;code&gt;defaultdict&lt;/code&gt; as a mapping but reconstructs a plain &lt;code&gt;dict&lt;/code&gt;, it is wrong, not merely imprecise.&lt;/p&gt;

&lt;h2&gt;
  
  
  Custom Containers: Either Register Them Or Treat Them As Leaves
&lt;/h2&gt;

&lt;p&gt;JAX and PyTorch are conservative about arbitrary subclasses. TensorFlow is more eager to recurse into sequence and mapping subclasses.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="k"&gt;class&lt;/span&gt; &lt;span class="nc"&gt;MyDict&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nb"&gt;dict&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;pass&lt;/span&gt;

&lt;span class="k"&gt;class&lt;/span&gt; &lt;span class="nc"&gt;MyList&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nb"&gt;list&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;pass&lt;/span&gt;

&lt;span class="k"&gt;class&lt;/span&gt; &lt;span class="nc"&gt;MyTuple&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nb"&gt;tuple&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;pass&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;JAX and PyTorch treat these as leaves unless explicitly registered:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nc"&gt;MyDict&lt;/span&gt;&lt;span class="p"&gt;({&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;}))[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# [MyDict({"b": 2, "a": 1})]
&lt;/span&gt;
&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nc"&gt;MyList&lt;/span&gt;&lt;span class="p"&gt;([&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;]))[&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;
&lt;span class="c1"&gt;# [MyList([1, 2])]
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;TensorFlow traverses them:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nc"&gt;MyDict&lt;/span&gt;&lt;span class="p"&gt;({&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;b&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;a&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;}))&lt;/span&gt;
&lt;span class="c1"&gt;# [1, 2]
&lt;/span&gt;
&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nc"&gt;MyList&lt;/span&gt;&lt;span class="p"&gt;([&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;]))&lt;/span&gt;
&lt;span class="c1"&gt;# [1, 2]
&lt;/span&gt;
&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;flatten&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nc"&gt;MyTuple&lt;/span&gt;&lt;span class="p"&gt;((&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;)))&lt;/span&gt;
&lt;span class="c1"&gt;# [1, 2]
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Namedtuple is the standard exception. All three frameworks recognize it as a structural container and preserve its type. They are also strict about namedtuple type matching: &lt;code&gt;Point(1, 2)&lt;/code&gt; is not the same structure as &lt;code&gt;(1, 2)&lt;/code&gt; or &lt;code&gt;RGB(1, 2)&lt;/code&gt;.&lt;/p&gt;

&lt;h2&gt;
  
  
  The Real Trap: &lt;code&gt;tree_map&lt;/code&gt; Does Not Always Mean Same-Structure Map
&lt;/h2&gt;

&lt;p&gt;JAX &lt;code&gt;tree_map&lt;/code&gt; uses the first argument as the reference structure. Later arguments are flattened "up to" that structure. If the first tree has a leaf, the corresponding value in a later tree may be an entire subtree.&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt; &lt;span class="p"&gt;[[&lt;/span&gt;&lt;span class="mi"&gt;3&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;4&lt;/span&gt;&lt;span class="p"&gt;}])&lt;/span&gt;
&lt;span class="c1"&gt;# [(1, [3]), (2, {"x": 4})]
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;The first tree says: "I am a list of two leaves." Therefore the second tree only needs to be a list of two objects. Those objects are passed whole to the function.&lt;/p&gt;

&lt;p&gt;The scalar case is even clearer:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;3&lt;/span&gt;&lt;span class="p"&gt;])&lt;/span&gt;
&lt;span class="c1"&gt;# (1, [2, 3])
&lt;/span&gt;
&lt;span class="n"&gt;jtu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt; &lt;span class="mi"&gt;3&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# ValueError: Expected list, got 3.
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;PyTorch 2.12.0 behaves similarly:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt; &lt;span class="p"&gt;[[&lt;/span&gt;&lt;span class="mi"&gt;3&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;4&lt;/span&gt;&lt;span class="p"&gt;}])&lt;/span&gt;
&lt;span class="c1"&gt;# [(1, [3]), (2, {"x": 4})]
&lt;/span&gt;
&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;3&lt;/span&gt;&lt;span class="p"&gt;])&lt;/span&gt;
&lt;span class="c1"&gt;# (1, [2, 3])
&lt;/span&gt;
&lt;span class="n"&gt;tpu&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tree_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt; &lt;span class="mi"&gt;3&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# ValueError: Node type mismatch; expected &amp;lt;class 'list'&amp;gt;, but got &amp;lt;class 'int'&amp;gt;.
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;PyTorch 2.2.2 did not support this multi-pytree call through &lt;code&gt;_pytree.tree_map&lt;/code&gt;. TensorFlow supports multiple structures, but it requires strict structural equality:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;nest&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;map_structure&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="k"&gt;lambda&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt; &lt;span class="p"&gt;[[&lt;/span&gt;&lt;span class="mi"&gt;3&lt;/span&gt;&lt;span class="p"&gt;],&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="mi"&gt;4&lt;/span&gt;&lt;span class="p"&gt;}])&lt;/span&gt;
&lt;span class="c1"&gt;# ValueError: structures do not have the same nested structure
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;h2&gt;
  
  
  Transform APIs: PyTree Support Is Not Just Flattening
&lt;/h2&gt;

&lt;p&gt;Tree semantics matter most when they meet transforms. Here the frameworks differ again.&lt;/p&gt;

&lt;p&gt;JAX transformations are natively pytree-based. &lt;code&gt;grad&lt;/code&gt; accepts nested inputs and returns gradients with the same structure:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="kn"&gt;import&lt;/span&gt; &lt;span class="n"&gt;jax&lt;/span&gt;
&lt;span class="kn"&gt;import&lt;/span&gt; &lt;span class="n"&gt;jax.numpy&lt;/span&gt; &lt;span class="k"&gt;as&lt;/span&gt; &lt;span class="n"&gt;jnp&lt;/span&gt;

&lt;span class="k"&gt;def&lt;/span&gt; &lt;span class="nf"&gt;f&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;**&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;][&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;**&lt;/span&gt; &lt;span class="mi"&gt;3&lt;/span&gt;

&lt;span class="n"&gt;params&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;jnp&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;array&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;2.0&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;jnp&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;array&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;)]}&lt;/span&gt;
&lt;span class="n"&gt;jax&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;grad&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;f&lt;/span&gt;&lt;span class="p"&gt;)(&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# {"x": Array(4., dtype=float32), "y": [Array(27., dtype=float32)]}
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;JAX &lt;code&gt;vmap&lt;/code&gt; accepts nested pytree inputs too:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="k"&gt;def&lt;/span&gt; &lt;span class="nf"&gt;g&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;][&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;

&lt;span class="n"&gt;batched&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;jnp&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;arange&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;jnp&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;arange&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;]}&lt;/span&gt;
&lt;span class="n"&gt;jax&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;vmap&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;g&lt;/span&gt;&lt;span class="p"&gt;)(&lt;/span&gt;&lt;span class="n"&gt;batched&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# Array([10., 12., 14.], dtype=float32)
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Because &lt;code&gt;None&lt;/code&gt; is a zero-leaf node in JAX, it can sit inside a vmapped input without becoming a batched argument:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="k"&gt;def&lt;/span&gt; &lt;span class="nf"&gt;h&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;

&lt;span class="n"&gt;jax&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;vmap&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;h&lt;/span&gt;&lt;span class="p"&gt;)({&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;jnp&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;arange&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="bp"&gt;None&lt;/span&gt;&lt;span class="p"&gt;]})&lt;/span&gt;
&lt;span class="c1"&gt;# Array([0., 1., 2.], dtype=float32)
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;Classic PyTorch autograd is different. &lt;code&gt;torch.autograd.grad&lt;/code&gt; expects tensors or gradient edges as &lt;code&gt;inputs&lt;/code&gt;, not an arbitrary nested pytree:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tensor&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;2.0&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;requires_grad&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="bp"&gt;True&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;y&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tensor&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;requires_grad&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="bp"&gt;True&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;loss&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;**&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt; &lt;span class="o"&gt;**&lt;/span&gt; &lt;span class="mi"&gt;3&lt;/span&gt;

&lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;autograd&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;grad&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;loss&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt;
&lt;span class="c1"&gt;# (tensor(4.), tensor(27.))
&lt;/span&gt;
&lt;span class="n"&gt;nested&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;]}&lt;/span&gt;
&lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;autograd&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;grad&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;loss&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;nested&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# RuntimeError: all inputs have to be Tensors or GradientEdges, but got str
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;The newer &lt;code&gt;torch.func&lt;/code&gt; stack does understand nested pytree-like parameter structures:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="kn"&gt;from&lt;/span&gt; &lt;span class="n"&gt;torch.func&lt;/span&gt; &lt;span class="kn"&gt;import&lt;/span&gt; &lt;span class="n"&gt;grad&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;vmap&lt;/span&gt;

&lt;span class="k"&gt;def&lt;/span&gt; &lt;span class="nf"&gt;f&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;**&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;][&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;**&lt;/span&gt; &lt;span class="mi"&gt;3&lt;/span&gt;

&lt;span class="n"&gt;params&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tensor&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;2.0&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;tensor&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;)]}&lt;/span&gt;
&lt;span class="nf"&gt;grad&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;f&lt;/span&gt;&lt;span class="p"&gt;)(&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# {"x": tensor(4.), "y": [tensor(27.)]}
&lt;/span&gt;
&lt;span class="k"&gt;def&lt;/span&gt; &lt;span class="nf"&gt;g&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;][&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;

&lt;span class="n"&gt;batched&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;arange&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;torch&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;arange&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;]}&lt;/span&gt;
&lt;span class="nf"&gt;vmap&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;g&lt;/span&gt;&lt;span class="p"&gt;)(&lt;/span&gt;&lt;span class="n"&gt;batched&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# tensor([10., 12., 14.])
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;TensorFlow's transform support follows &lt;code&gt;tf.nest&lt;/code&gt;. &lt;code&gt;GradientTape.gradient&lt;/code&gt; accepts nested sources and returns gradients in the same structure:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;x&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nc"&gt;Variable&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;2.0&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;y&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nc"&gt;Variable&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;nested&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;x&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;y&lt;/span&gt;&lt;span class="p"&gt;]}&lt;/span&gt;

&lt;span class="k"&gt;with&lt;/span&gt; &lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nc"&gt;GradientTape&lt;/span&gt;&lt;span class="p"&gt;()&lt;/span&gt; &lt;span class="k"&gt;as&lt;/span&gt; &lt;span class="n"&gt;tape&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt;
    &lt;span class="n"&gt;loss&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;nested&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;**&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;nested&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;][&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;**&lt;/span&gt; &lt;span class="mi"&gt;3&lt;/span&gt;

&lt;span class="n"&gt;tape&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;gradient&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;loss&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;nested&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# {"x": tf.Tensor(4.0), "y": [tf.Tensor(27.0)]}
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;&lt;code&gt;tf.vectorized_map&lt;/code&gt; also accepts nested input structures:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="k"&gt;def&lt;/span&gt; &lt;span class="nf"&gt;g&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;][&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt;

&lt;span class="n"&gt;batched&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="p"&gt;{&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;range&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;range&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="mi"&gt;10&lt;/span&gt;&lt;span class="p"&gt;]}&lt;/span&gt;
&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;vectorized_map&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;g&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;batched&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="c1"&gt;# tf.Tensor([10. 12. 14.], shape=(3,), dtype=float32)
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;&lt;code&gt;tf.function&lt;/code&gt; accepts nested structures as ordinary function arguments:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="nd"&gt;@tf.function&lt;/span&gt;
&lt;span class="k"&gt;def&lt;/span&gt; &lt;span class="nf"&gt;f&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;return&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;**&lt;/span&gt; &lt;span class="mi"&gt;2&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;][&lt;/span&gt;&lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;]&lt;/span&gt; &lt;span class="o"&gt;**&lt;/span&gt; &lt;span class="mi"&gt;3&lt;/span&gt;

&lt;span class="nf"&gt;f&lt;/span&gt;&lt;span class="p"&gt;({&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;x&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;constant&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;2.0&lt;/span&gt;&lt;span class="p"&gt;),&lt;/span&gt; &lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="s"&gt;y&lt;/span&gt;&lt;span class="sh"&gt;"&lt;/span&gt;&lt;span class="p"&gt;:&lt;/span&gt; &lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;tf&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;constant&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="mf"&gt;3.0&lt;/span&gt;&lt;span class="p"&gt;)]})&lt;/span&gt;
&lt;span class="c1"&gt;# tf.Tensor(31.0, shape=(), dtype=float32)
&lt;/span&gt;&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;The right summary is more specific: JAX transforms are pytree-native; PyTorch classic autograd is not, while &lt;code&gt;torch.func&lt;/code&gt; is; TensorFlow transform APIs accept nested structures.&lt;/p&gt;

&lt;h2&gt;
  
  
  Closing
&lt;/h2&gt;

&lt;p&gt;PyTrees are a small abstraction with a long tail. Simple examples make every framework look compatible; real optimizer states, optional values, ordered mappings, custom containers, and transform APIs expose the differences quickly.&lt;/p&gt;

</description>
      <category>jax</category>
      <category>torch</category>
      <category>tensorflow</category>
    </item>
    <item>
      <title>TensorCircuit-NG vs cuQuantum on H200: JIT compilation beats the "magic GPU library" assumption</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Sun, 07 Jun 2026 02:02:29 +0000</pubDate>
      <link>https://dev.to/refractionray/tensorcircuit-ng-vs-cuquantum-on-h200-jit-compilation-beats-the-magic-gpu-library-assumption-d5c</link>
      <guid>https://dev.to/refractionray/tensorcircuit-ng-vs-cuquantum-on-h200-jit-compilation-beats-the-magic-gpu-library-assumption-d5c</guid>
      <description>&lt;p&gt;NVIDIA cuQuantum has a strong reputation as the natural high-performance baseline for GPU quantum simulation. That reputation is understandable: cuQuantum contains serious low-level GPU libraries such as cuStateVec and cuTensorNet and it is NVIDIA who creates GPU and CUDA!&lt;/p&gt;

&lt;p&gt;But in an end-to-end differentiable VQE workload, the result is more nuanced. On our H200 GPU benchmark, TensorCircuit-NG was substantially faster after compilation, while also offering a much higher-level and user-friendly programming model.&lt;/p&gt;

&lt;p&gt;The short version:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;cuQuantum is a powerful low-level library.&lt;/li&gt;
&lt;li&gt;It is not automatically the fastest route for practical quantum simulation tasks.&lt;/li&gt;
&lt;li&gt;Direct cuQuantum code is significantly more verbose and engineering-heavy.&lt;/li&gt;
&lt;li&gt;TensorCircuit-NG pays a JAX compilation cost, but repeated value-and-gradient evaluations quickly amortize that cost.&lt;/li&gt;
&lt;li&gt;The final running time of TensorCircuit-NG is much shorter than NVIDIA cuquantum.&lt;/li&gt;
&lt;/ul&gt;

&lt;h2&gt;
  
  
  Benchmark setup
&lt;/h2&gt;

&lt;p&gt;We used the workload as in &lt;a href="https://github.com/tensorcircuit/tensorcircuit-ng/blob/master/examples/benchmark_cuquantum_vs_tc_vqe.py" rel="noopener noreferrer"&gt;the script&lt;/a&gt; for 1D TFIM VQE task:&lt;/p&gt;

&lt;p&gt;Hardware and software:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;GPU: NVIDIA H200&lt;/li&gt;
&lt;li&gt;TensorCircuit-NG: &lt;code&gt;1.6.0&lt;/code&gt;
&lt;/li&gt;
&lt;li&gt;JAX: &lt;code&gt;0.7.2&lt;/code&gt;
&lt;/li&gt;
&lt;li&gt;cuQuantum Python: &lt;code&gt;26.3.2&lt;/code&gt;
&lt;/li&gt;
&lt;li&gt;CuPy: &lt;code&gt;14.1.1&lt;/code&gt;
&lt;/li&gt;
&lt;li&gt;PyTorch: &lt;code&gt;2.11.0+cu128&lt;/code&gt;
&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;We measured one warmup/compile call and then the mean of five later value-and-gradient calls.&lt;/p&gt;

&lt;h2&gt;
  
  
  Implementations compared
&lt;/h2&gt;

&lt;p&gt;We tested two TensorCircuit-NG modes:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;
&lt;strong&gt;TC-JAX scan&lt;/strong&gt;: uses &lt;code&gt;scan&lt;/code&gt; over VQE layers to reduce JAX compilation/staging time.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;TC-JAX unrolled&lt;/strong&gt;: builds all layers directly. This produces a larger traced program, but can be faster after compilation.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;We also tested two direct cuQuantum routes:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;
&lt;strong&gt;cuStateVec adjoint&lt;/strong&gt;: applies gates with cuStateVec and computes the full gradient with adjoint differentiation. This is not parameter shift so it is a fair comparison.&lt;/li&gt;
&lt;li&gt;
&lt;strong&gt;cuTensorNet full-state autograd&lt;/strong&gt;: contracts the full state with cuTensorNet, then computes the TFIM state-vector expectation on GPU with PyTorch autograd.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;The cuTensorNet path is intentionally not the obviously bad version where every Pauli term gets a separate tensor-network path search. We first tried that more "TN-native" observable-contraction style, but for this workload it spent too much time in repeated graph/path overhead. The final version is closer to the state-vector expectation workflow used by the TensorCircuit-NG and MindQuantum benchmark.&lt;/p&gt;

&lt;h2&gt;
  
  
  Repeated value-and-gradient runtime
&lt;/h2&gt;

&lt;p&gt;The table below reports the post-warmup runtime. This is the relevant metric for VQE-style optimization, where the same circuit structure is evaluated many times.&lt;/p&gt;

&lt;div class="table-wrapper-paragraph"&gt;&lt;table&gt;
&lt;thead&gt;
&lt;tr&gt;
&lt;th&gt;backend&lt;/th&gt;
&lt;th&gt;14 qubits&lt;/th&gt;
&lt;th&gt;20 qubits&lt;/th&gt;
&lt;th&gt;24 qubits&lt;/th&gt;
&lt;/tr&gt;
&lt;/thead&gt;
&lt;tbody&gt;
&lt;tr&gt;
&lt;td&gt;TC-JAX scan&lt;/td&gt;
&lt;td&gt;0.01201s&lt;/td&gt;
&lt;td&gt;0.01616s&lt;/td&gt;
&lt;td&gt;0.06374s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;TC-JAX unrolled&lt;/td&gt;
&lt;td&gt;0.00995s&lt;/td&gt;
&lt;td&gt;0.01381s&lt;/td&gt;
&lt;td&gt;0.02547s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;cuStateVec adjoint&lt;/td&gt;
&lt;td&gt;0.08036s&lt;/td&gt;
&lt;td&gt;0.12061s&lt;/td&gt;
&lt;td&gt;0.30142s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;cuTensorNet full-state autograd&lt;/td&gt;
&lt;td&gt;1.35677s&lt;/td&gt;
&lt;td&gt;2.04291s&lt;/td&gt;
&lt;td&gt;2.30414s&lt;/td&gt;
&lt;/tr&gt;
&lt;/tbody&gt;
&lt;/table&gt;&lt;/div&gt;

&lt;p&gt;In repeated value-and-gradient calls, TensorCircuit-NG is faster than cuStateVec:&lt;/p&gt;

&lt;div class="table-wrapper-paragraph"&gt;&lt;table&gt;
&lt;thead&gt;
&lt;tr&gt;
&lt;th&gt;qubits&lt;/th&gt;
&lt;th&gt;TC-JAX scan vs cuStateVec&lt;/th&gt;
&lt;th&gt;TC-JAX unrolled vs cuStateVec&lt;/th&gt;
&lt;/tr&gt;
&lt;/thead&gt;
&lt;tbody&gt;
&lt;tr&gt;
&lt;td&gt;14&lt;/td&gt;
&lt;td&gt;6.69x&lt;/td&gt;
&lt;td&gt;8.08x&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;20&lt;/td&gt;
&lt;td&gt;7.46x&lt;/td&gt;
&lt;td&gt;8.73x&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;24&lt;/td&gt;
&lt;td&gt;4.73x&lt;/td&gt;
&lt;td&gt;11.83x&lt;/td&gt;
&lt;/tr&gt;
&lt;/tbody&gt;
&lt;/table&gt;&lt;/div&gt;

&lt;p&gt;The gap is much larger against the cuTensorNet route for this particular state-vector expectation plus autograd workflow:&lt;/p&gt;

&lt;div class="table-wrapper-paragraph"&gt;&lt;table&gt;
&lt;thead&gt;
&lt;tr&gt;
&lt;th&gt;qubits&lt;/th&gt;
&lt;th&gt;TC-JAX scan vs cuTensorNet&lt;/th&gt;
&lt;th&gt;TC-JAX unrolled vs cuTensorNet&lt;/th&gt;
&lt;/tr&gt;
&lt;/thead&gt;
&lt;tbody&gt;
&lt;tr&gt;
&lt;td&gt;14&lt;/td&gt;
&lt;td&gt;112.97x&lt;/td&gt;
&lt;td&gt;136.36x&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;20&lt;/td&gt;
&lt;td&gt;126.42x&lt;/td&gt;
&lt;td&gt;147.93x&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;24&lt;/td&gt;
&lt;td&gt;36.15x&lt;/td&gt;
&lt;td&gt;90.46x&lt;/td&gt;
&lt;/tr&gt;
&lt;/tbody&gt;
&lt;/table&gt;&lt;/div&gt;

&lt;p&gt;These numbers are the main point: cuQuantum is not a magic speed button. A library being close to CUDA, or being written by a GPU vendor, does not automatically make it the fastest end-to-end implementation for a differentiable quantum algorithm.&lt;/p&gt;

&lt;h2&gt;
  
  
  First-call cost and amortization
&lt;/h2&gt;

&lt;p&gt;cuQuantum has much lower first-call overhead. This is expected: TensorCircuit-NG uses JAX JIT compilation, and that first call can be expensive.&lt;/p&gt;

&lt;p&gt;So if the task is a single one-off circuit evaluation, cuQuantum's low startup cost is attractive. But VQE is usually not a one-off workload. It repeatedly evaluates the same circuit structure for many optimizer steps and often across multiple random initializations. In that regime, TensorCircuit-NG's first-call cost is easily amortized, and the much faster post-compilation runtime becomes the dominant factor.&lt;/p&gt;

&lt;p&gt;There is also a useful TensorCircuit-NG tradeoff:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;Use &lt;strong&gt;scan mode&lt;/strong&gt; when compilation time matters.&lt;/li&gt;
&lt;li&gt;Use &lt;strong&gt;unrolled mode&lt;/strong&gt; when the same circuit will be evaluated many times and peak post-compilation throughput matters.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;At 24 qubits, unrolled TensorCircuit-NG is about &lt;code&gt;2.50x&lt;/code&gt; faster than scan mode after compilation, but the first call is about &lt;code&gt;9x&lt;/code&gt; heavier.&lt;/p&gt;

&lt;h2&gt;
  
  
  Programming model
&lt;/h2&gt;

&lt;p&gt;Performance is only half of the story. The programming model matters.&lt;/p&gt;

&lt;p&gt;In TensorCircuit-NG, the benchmark is expressed as circuit code:&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight python"&gt;&lt;code&gt;&lt;span class="n"&gt;c&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tc&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nc"&gt;Circuit&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;n&lt;/span&gt;&lt;span class="p"&gt;)&lt;/span&gt;
&lt;span class="n"&gt;c&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;h&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="nf"&gt;range&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;n&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt;
&lt;span class="k"&gt;for&lt;/span&gt; &lt;span class="n"&gt;layer&lt;/span&gt; &lt;span class="ow"&gt;in&lt;/span&gt; &lt;span class="nf"&gt;range&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;depth&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
    &lt;span class="k"&gt;for&lt;/span&gt; &lt;span class="n"&gt;i&lt;/span&gt; &lt;span class="ow"&gt;in&lt;/span&gt; &lt;span class="nf"&gt;range&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;n&lt;/span&gt; &lt;span class="o"&gt;-&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
        &lt;span class="n"&gt;c&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;rzz&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;i&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;i&lt;/span&gt; &lt;span class="o"&gt;+&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;theta&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;layer&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;0&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;i&lt;/span&gt;&lt;span class="p"&gt;])&lt;/span&gt;
    &lt;span class="k"&gt;for&lt;/span&gt; &lt;span class="n"&gt;i&lt;/span&gt; &lt;span class="ow"&gt;in&lt;/span&gt; &lt;span class="nf"&gt;range&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;n&lt;/span&gt;&lt;span class="p"&gt;):&lt;/span&gt;
        &lt;span class="n"&gt;c&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;rx&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;i&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;theta&lt;/span&gt;&lt;span class="o"&gt;=&lt;/span&gt;&lt;span class="n"&gt;params&lt;/span&gt;&lt;span class="p"&gt;[&lt;/span&gt;&lt;span class="n"&gt;layer&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="mi"&gt;1&lt;/span&gt;&lt;span class="p"&gt;,&lt;/span&gt; &lt;span class="n"&gt;i&lt;/span&gt;&lt;span class="p"&gt;])&lt;/span&gt;

&lt;span class="n"&gt;value_and_grad&lt;/span&gt; &lt;span class="o"&gt;=&lt;/span&gt; &lt;span class="n"&gt;tc&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;backend&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;jit&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;tc&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="n"&gt;backend&lt;/span&gt;&lt;span class="p"&gt;.&lt;/span&gt;&lt;span class="nf"&gt;value_and_grad&lt;/span&gt;&lt;span class="p"&gt;(&lt;/span&gt;&lt;span class="n"&gt;energy_fn&lt;/span&gt;&lt;span class="p"&gt;))&lt;/span&gt;
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;With direct cuQuantum, the user has to manually manage much lower-level details:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;gate matrices and their dtype conventions&lt;/li&gt;
&lt;li&gt;state-vector memory&lt;/li&gt;
&lt;li&gt;cuStateVec binding signatures&lt;/li&gt;
&lt;li&gt;tensor-network modes&lt;/li&gt;
&lt;li&gt;PyTorch operands for autograd&lt;/li&gt;
&lt;li&gt;GPU synchronization&lt;/li&gt;
&lt;li&gt;version-specific API behavior&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;cuQuantum is valuable, but it is closer to a low-level engine than a high-level quantum algorithm framework. For a researcher, that difference is very real.&lt;/p&gt;

&lt;h2&gt;
  
  
  Takeaway
&lt;/h2&gt;

&lt;p&gt;This benchmark does &lt;strong&gt;not&lt;/strong&gt; prove that cuQuantum is slow for every task.  What this benchmark does show is narrower and more practical:&lt;/p&gt;

&lt;p&gt;For VQE workload, direct cuQuantum was not the fastest end-to-end route. TensorCircuit-NG provided a much simpler programming interface and substantially faster repeated value-and-gradient evaluations after JAX compilation.&lt;/p&gt;

&lt;p&gt;The common assumption that "NVIDIA controls CUDA, therefore cuQuantum must be the fastest implementation" is too simplistic. Raw GPU kernels matter, but so do JIT compilation, autodiff integration, graph-level optimization, and the abstraction level exposed to users.&lt;/p&gt;

&lt;p&gt;TensorCircuit-NG's advantage is that it lets users write concise quantum-program code while still compiling to high-performance backend-native tensor programs. For repeated VQE-style workloads, that combination can beat direct cuQuantum both in usability and in runtime.&lt;/p&gt;

</description>
      <category>python</category>
      <category>gpu</category>
      <category>cuda</category>
    </item>
    <item>
      <title>Why JAX Is a Much Better Backend for Quantum Circuit Simulation Than PyTorch</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Sat, 06 Jun 2026 05:01:36 +0000</pubDate>
      <link>https://dev.to/refractionray/why-jax-is-a-much-better-backend-for-quantum-circuit-simulation-than-pytorch-ak6</link>
      <guid>https://dev.to/refractionray/why-jax-is-a-much-better-backend-for-quantum-circuit-simulation-than-pytorch-ak6</guid>
      <description>&lt;p&gt;Modern quantum circuit simulation is not just “machine learning with complex tensors.” It involves irregular tensor contractions, sparse operators, statevector transformations, and automatic differentiation through all of them. This makes backend choice unusually important. A backend that is excellent for standard neural-network layers may still be a poor fit for general quantum simulation workloads.&lt;/p&gt;

&lt;p&gt;We benchmarked this with a simple VQE workload for the 1D transverse-field Ising&lt;br&gt;
model as in &lt;a href="https://github.com/tensorcircuit/tensorcircuit-ng/blob/master/examples/benchmark_jax_vs_torch_vqe.py" rel="noopener noreferrer"&gt;the script&lt;/a&gt;,&lt;br&gt;
&lt;/p&gt;

&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight plaintext"&gt;&lt;code&gt;H = -sum_i Z_i Z_{i+1} - sum_i X_i,
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;using 20 qubits, 10 ansatz layers, complex64 precision, and one NVIDIA RTX 5090 GPU. &lt;/p&gt;

&lt;h2&gt;
  
  
  Results
&lt;/h2&gt;

&lt;div class="table-wrapper-paragraph"&gt;&lt;table&gt;
&lt;thead&gt;
&lt;tr&gt;
&lt;th&gt;Backend&lt;/th&gt;
&lt;th&gt;Compile / Warmup&lt;/th&gt;
&lt;th&gt;Value+Grad Runtime&lt;/th&gt;
&lt;/tr&gt;
&lt;/thead&gt;
&lt;tbody&gt;
&lt;tr&gt;
&lt;td&gt;TensorCircuit-NG, JAX backend&lt;/td&gt;
&lt;td&gt;53.53 s&lt;/td&gt;
&lt;td&gt;0.0265 s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;TensorCircuit-NG, PyTorch backend&lt;/td&gt;
&lt;td&gt;0.48 s&lt;/td&gt;
&lt;td&gt;0.3299 s&lt;/td&gt;
&lt;/tr&gt;
&lt;tr&gt;
&lt;td&gt;TorchQuantum, optimized implementation than default&lt;/td&gt;
&lt;td&gt;0.81 s&lt;/td&gt;
&lt;td&gt;0.4172 s&lt;/td&gt;
&lt;/tr&gt;
&lt;/tbody&gt;
&lt;/table&gt;&lt;/div&gt;

&lt;p&gt;The JAX backend is about &lt;strong&gt;12.4x faster&lt;/strong&gt; than TensorCircuit-NG’s PyTorch backend and about &lt;strong&gt;15.7x faster&lt;/strong&gt; than TorchQuantum for the post-compilation value-and-gradient step.&lt;/p&gt;

&lt;p&gt;The compile time tells the other half of the story: JAX pays a much larger upfront XLA compilation cost. But after compilation, XLA produces a far more effective execution plan for this quantum simulation workload. This is exactly the tradeoff we want in VQE, QAOA, time evolution, and many other iterative algorithms: pay once, run many times.&lt;/p&gt;

&lt;h2&gt;
  
  
  Why This Happens
&lt;/h2&gt;

&lt;p&gt;Quantum circuit simulation stresses a backend differently from ordinary deep learning. The workload mixes tensor-network contraction, sparse Hamiltonian application, and reverse-mode differentiation. JAX/XLA is designed to see the whole computation and optimize it aggressively as a compiled program on the target device.&lt;/p&gt;

&lt;p&gt;PyTorch, in contrast, is strongest where the workload resembles standard neural network layers. For more general tensor programs, especially tensor-network-like simulation code, the compiler stack is less aggressive and less predictable.&lt;br&gt;
In this benchmark, the same TensorCircuit-NG algorithm is more than an order of magnitude faster on JAX than on PyTorch after compilation.&lt;/p&gt;

&lt;h2&gt;
  
  
  A Note on TorchQuantum
&lt;/h2&gt;

&lt;p&gt;We also compared against TorchQuantum as a representative PyTorch-native quantum circuit package. To make the comparison generous, we did not use its generic Pauli-string expectation path. That built-in route tends to materialize dense Pauli operators and is slow and not scalable. Instead, we implemented a TFIM-specific expectation directly extracted from state:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;
&lt;code&gt;ZZ&lt;/code&gt; terms are evaluated from probabilities and precomputed sign tensors.&lt;/li&gt;
&lt;li&gt;
&lt;code&gt;X&lt;/code&gt; terms are evaluated by flipping the state axis and taking an inner product.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;This is already a substantial low-level optimization Even with that help, TorchQuantum remains slower than TensorCircuit-NG on the JAX backend by about 15.7x. And even if you prefer PyTorch backend, PyTorch backend from TensorCircuit-NG is still a better choice in terms of both warm-up and run times.&lt;/p&gt;

&lt;h2&gt;
  
  
  Takeaway
&lt;/h2&gt;

&lt;p&gt;The lesson is not merely that one package is faster than another. The deeper point is that backend architecture matters. Quantum simulation benefits from a compiler that can optimize a whole differentiable tensor program, not just a collection of familiar machine-learning layers.&lt;/p&gt;

&lt;p&gt;For TensorCircuit-NG, the JAX backend gives exactly that: a high-level quantum programming interface backed by XLA’s aggressive compilation. The result is a backend that is not only elegant for research code, but also dramatically faster for real differentiable quantum simulation workloads.&lt;/p&gt;

</description>
      <category>jax</category>
      <category>pytorch</category>
    </item>
    <item>
      <title>TensorCircuit-NG: How to Tell Whether a Quantum x AI x HPC Platform Is Truly Mature When Everyone Tells the Same Story</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Thu, 04 Jun 2026 08:28:33 +0000</pubDate>
      <link>https://dev.to/refractionray/tensorcircuit-ng-how-to-tell-whether-a-quantum-x-ai-x-hpc-platform-is-truly-mature-when-everyone-577p</link>
      <guid>https://dev.to/refractionray/tensorcircuit-ng-how-to-tell-whether-a-quantum-x-ai-x-hpc-platform-is-truly-mature-when-everyone-577p</guid>
      <description>&lt;p&gt;In recent years, the convergence of quantum computing, artificial intelligence (AI), and high-performance computing (HPC) has become a central theme in the evolution of scientific computing infrastructure. From AI4Science and quantum machine learning to supercomputing centers and heterogeneous computing platforms, phrases such as "Quantum x AI x HPC", "integrated quantum-supercomputing-intelligence infrastructure", and "next-generation research infrastructure" now appear frequently in academic conferences, industry forums, and corporate presentations.&lt;/p&gt;

&lt;p&gt;At the same time, a clear pattern has emerged:&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;The messaging is becoming increasingly similar, while the actual technical depth of different products varies dramatically.&lt;/strong&gt;&lt;/p&gt;

&lt;p&gt;Whether the subject is a quantum software platform, an AI4Science infrastructure stack, or a heterogeneous computing framework, many projects now describe themselves in similar terms:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;integrating quantum computing with AI;&lt;/li&gt;
&lt;li&gt;supporting heterogeneous computing resources;&lt;/li&gt;
&lt;li&gt;serving as future research infrastructure;&lt;/li&gt;
&lt;li&gt;enabling applications in materials science, chemistry, biomedicine, and other industries;&lt;/li&gt;
&lt;li&gt;building an open ecosystem and developer community.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;These directions are meaningful. In fact, they are becoming part of the field's shared consensus.&lt;/p&gt;

&lt;p&gt;The real question is different:&lt;/p&gt;

&lt;p&gt;&lt;strong&gt;When everyone is telling a similar story, how can we tell whether a platform has actually delivered technical substance, rather than remaining at the level of conceptual packaging and slideware?&lt;/strong&gt;&lt;/p&gt;

&lt;p&gt;For scientific infrastructure, it is more useful to ask five verifiable questions than to focus on slogans:&lt;/p&gt;

&lt;ol&gt;
&lt;li&gt;Is it open source?&lt;/li&gt;
&lt;li&gt;Does it provide public benchmarks?&lt;/li&gt;
&lt;li&gt;Is it used continuously by high-quality research communities?&lt;/li&gt;
&lt;li&gt;Has it supported real industry-oriented application cases?&lt;/li&gt;
&lt;li&gt;Does it continue to evolve through sustained version updates?&lt;/li&gt;
&lt;/ol&gt;

&lt;p&gt;Any platform that claims to be "next-generation research infrastructure" should be able to answer these questions in a concrete way.&lt;/p&gt;

&lt;p&gt;The development of &lt;a href="https://github.com/tensorcircuit/tensorcircuit-ng" rel="noopener noreferrer"&gt;TensorCircuit-NG&lt;/a&gt; offers a useful case study. Its value does not lie only in proposing a vision for "Quantum x AI x HPC"; it lies in a body of work that can be inspected, reproduced, cited, extended, and tested over time: open code, reproducible performance evaluations, a visible record of academic adoption, evidence of industry spillover, and six years of engineering iteration.&lt;/p&gt;

&lt;h2&gt;
  
  
  1. Is It Open Source?
&lt;/h2&gt;

&lt;p&gt;For scientific software, open source means more than publishing code.&lt;/p&gt;

&lt;p&gt;It means that:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;the technology can be independently verified;&lt;/li&gt;
&lt;li&gt;performance claims can be reproduced;&lt;/li&gt;
&lt;li&gt;algorithms can be inspected;&lt;/li&gt;
&lt;li&gt;users can deploy the software without relying on a closed service;&lt;/li&gt;
&lt;li&gt;third-party researchers can repeat experiments under their own conditions.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;Research communities do not lack polished presentations. What is much rarer is a technical system that can survive independent inspection.&lt;/p&gt;

&lt;p&gt;TensorCircuit was not released as a one-off code dump. Its development forms a traceable engineering trajectory: from the original personal open-source version, to the version developed during the Tencent Quantum Lab period, and then to the currently maintained TensorCircuit-NG project. Across these stages, the core code, documentation, tests, and examples have remained open. The GitHub history preserves the development record, with more than 500 combined stars and forks, over 2,700 commits, more than 30 released versions, and contributions from over 30 developers around the world.&lt;/p&gt;

&lt;p&gt;In terms of engineering scale, TensorCircuit-NG is no longer a short-term proof-of-concept project. It is a platform-level scientific computing system, with roughly 70,000 lines of code, type annotations, unit tests, continuous integration, documentation, and tutorials. The repository currently contains close to one thousand test functions. These tests are not merely a coverage metric; they are part of the engineering foundation that keeps APIs stable, backend behavior consistent, and long-term maintenance manageable.&lt;/p&gt;

&lt;p&gt;The surrounding ecosystem matters as well. TensorCircuit-NG provides documentation, more than 30 tutorial examples, over 170 application examples, more than 10 benchmark suites, and a companion quantum computing tutorial. Together, these resources form a developer ecosystem that is learnable, reusable, and extensible. The platform also embraces AI-native workflows by providing AI skill packages for paper reproduction, code translation, and performance optimization. This means TensorCircuit-NG is not only designed for human developers; it is also adapting to a new mode of scientific software development in which AI agents participate directly in research workflows.&lt;/p&gt;

&lt;p&gt;Another measurable signal of open-source adoption is installation and use. TensorCircuit-related packages include &lt;code&gt;tensorcircuit&lt;/code&gt;, &lt;code&gt;tensorcircuit-ng&lt;/code&gt;, and the nightly package &lt;code&gt;tensorcircuit-nightly&lt;/code&gt; on PyPI, with cumulative &lt;code&gt;pip install&lt;/code&gt; downloads exceeding one million. Download counts alone do not prove scientific value, but they do show that the platform exists in real development environments, not only in papers or promotional pages.&lt;/p&gt;

&lt;p&gt;For research infrastructure, credibility comes from the ability of third-party users to run the code, inspect the implementation, reproduce experiments, and build their own workflows. Code is always more honest than marketing material.&lt;/p&gt;

&lt;h2&gt;
  
  
  2. Are There Public Benchmarks?
&lt;/h2&gt;

&lt;p&gt;Every computing platform eventually has to answer a simple question:&lt;/p&gt;

&lt;p&gt;Does it actually improve computational efficiency?&lt;/p&gt;

&lt;p&gt;This is why public benchmarking is essential for judging platform maturity. In a high-performance setting such as "Quantum x AI x HPC", claims about acceleration, heterogeneous execution, or scalability are difficult to evaluate without reproducible benchmarks.&lt;/p&gt;

&lt;p&gt;One of the earliest reasons TensorCircuit attracted attention was its benchmark system for differentiable quantum computing and tensor-network simulation. The first TensorCircuit white paper was published in &lt;em&gt;Quantum&lt;/em&gt;: &lt;a href="https://quantum-journal.org/papers/q-2023-02-02-912/" rel="noopener noreferrer"&gt;TensorCircuit: a Quantum Software Framework for the NISQ Era&lt;/a&gt;. The paper introduced the platform architecture, core functionality, and performance advantages, and compared TensorCircuit against several mainstream quantum software frameworks on variational quantum algorithms, gradient computation, and quantum circuit simulation.&lt;/p&gt;

&lt;p&gt;The work made clear why unified tensor programming, automatic differentiation, and just-in-time compilation matter for quantum computing workflows. In several variational quantum algorithm and gradient computation tasks, TensorCircuit demonstrated significant performance advantages over representative frameworks such as IBM's Qiskit and PennyLane, with speedups reaching multiple orders of magnitude in some cases. More importantly, these results were not confined to figures in a paper: the code, experimental setup, and evaluation procedures were made reproducible. That is the difference between a verifiable technical path and an unverifiable performance claim.&lt;/p&gt;

&lt;p&gt;With the release of TensorCircuit-NG, the benchmark scope has expanded toward problems closer to future research infrastructure:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;GPU-accelerated computing;&lt;/li&gt;
&lt;li&gt;optimized tensor-network contraction;&lt;/li&gt;
&lt;li&gt;distributed HPC environments;&lt;/li&gt;
&lt;li&gt;unified computation graphs spanning quantum circuits, neural networks, and tensor networks.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;The NG white paper further summarizes TensorCircuit's upgrade toward the integration of quantum computing, supercomputing, and intelligent computing; see the &lt;a href="https://arxiv.org/abs/2602.14167" rel="noopener noreferrer"&gt;preprint&lt;/a&gt;. The focus has shifted from "how to simulate quantum circuits faster on a single machine" to "how to organize quantum, AI, and numerical computing workflows in realistic heterogeneous research environments."&lt;/p&gt;

&lt;p&gt;External evaluation provides another layer of evidence. NVIDIA used TensorCircuit as a third-party quantum software case in its cuQuantum 23.10 benchmarking context. This shows that TensorCircuit has entered the evaluation landscape of hardware and high-performance computing vendors. For scientific infrastructure, such external benchmarks complement open papers and are more persuasive than slide-based claims.&lt;/p&gt;

&lt;h2&gt;
  
  
  3. Is It Used by the Research Community?
&lt;/h2&gt;

&lt;p&gt;For scientific infrastructure, the hardest signal to fake is not performance.&lt;/p&gt;

&lt;p&gt;It is sustained use by serious research communities.&lt;/p&gt;

&lt;p&gt;A platform can gain short-term attention through marketing, but it cannot gain long-term citations through marketing alone. Research adoption is a form of long-horizon voting. If a platform continues to support high-quality work across institutions, research areas, and teams, then it has demonstrated real utility.&lt;/p&gt;

&lt;p&gt;More than 170 academic works have cited TensorCircuit, and in the first five months of 2026 alone, more than 40 works have already cited it. More importantly, these works are not concentrated in a single niche. They span quantum simulation, quantum machine learning, quantum chemistry, quantum sensing, quantum architecture search, and AI4Science.&lt;/p&gt;

&lt;h3&gt;
  
  
  Quantum Simulation and Many-Body Systems
&lt;/h3&gt;

&lt;p&gt;In many-body quantum physics, condensed matter systems, and complex quantum dynamics, researchers often need large-scale quantum circuit simulation, tensor-network contraction, and differentiable optimization. These tasks place high demands on performance, numerical stability, and automatic differentiation.&lt;/p&gt;

&lt;p&gt;Representative works include &lt;a href="https://quantum-journal.org/papers/q-2024-07-23-1422/" rel="noopener noreferrer"&gt;Zero and Finite Temperature Quantum Simulations Powered by Quantum Magic&lt;/a&gt; from teams including NVIDIA, Google, MIT, and Harvard; &lt;a href="https://arxiv.org/abs/2501.04679/" rel="noopener noreferrer"&gt;Exploring nontrivial topology at quantum criticality in a superconducting processor&lt;/a&gt; from Haohua Wang's group at Zhejiang University; and &lt;a href="https://arxiv.org/abs/2409.07281" rel="noopener noreferrer"&gt;Variational LOCC-assisted quantum circuits for long-range entangled states&lt;/a&gt; from Xiongfeng Ma's group at Tsinghua University. These papers show that TensorCircuit is not limited to abstract algorithm demonstrations; it is being used in concrete problems in many-body physics and experimental quantum information.&lt;/p&gt;

&lt;h3&gt;
  
  
  Quantum Machine Learning
&lt;/h3&gt;

&lt;p&gt;Quantum machine learning is one of the most active application areas for TensorCircuit. Representative papers include &lt;a href="https://www.nature.com/articles/s41467-024-45882-z" rel="noopener noreferrer"&gt;Understanding quantum machine learning also requires rethinking generalization&lt;/a&gt; from Jens Eisert's group at the Free University of Berlin, &lt;a href="https://www.nature.com/articles/s41467-024-53769-2" rel="noopener noreferrer"&gt;Dynamical transition in controllable quantum neural networks with large depth&lt;/a&gt; from teams including Liang Jiang and Junyu Liu, &lt;a href="https://journals.aps.org/prl/abstract/10.1103/PhysRevLett.132.100602" rel="noopener noreferrer"&gt;Generative Quantum Machine Learning via Denoising Diffusion Probabilistic Models&lt;/a&gt; from Quntao Zhuang's group at the University of Southern California, and IBM Quantum's &lt;a href="https://arxiv.org/abs/2411.05760" rel="noopener noreferrer"&gt;Dynamic parameterized quantum circuits: expressive and barren-plateau free&lt;/a&gt;.&lt;/p&gt;

&lt;p&gt;These works all require stable workflows connecting parameterized quantum circuits, gradient computation, model training, and numerical simulation. TensorCircuit's value is visible precisely at this workflow level: it connects quantum circuit simulation, automatic differentiation, and machine learning training into a unified programmable system.&lt;/p&gt;

&lt;h3&gt;
  
  
  Quantum Architecture Search and Algorithm Design
&lt;/h3&gt;

&lt;p&gt;TensorCircuit has also been used in algorithmic and learning-theoretic research. Examples include &lt;a href="https://journals.aps.org/prxquantum/abstract/10.1103/PRXQuantum.5.040306" rel="noopener noreferrer"&gt;Learning Quantum States and Unitaries of Bounded Gate Complexity&lt;/a&gt; from Caltech and Google, &lt;a href="https://ieeexplore.ieee.org/abstract/document/10821373" rel="noopener noreferrer"&gt;Quantum Machine Learning Architecture Search via Deep Reinforcement Learning&lt;/a&gt; from Brookhaven National Laboratory, and &lt;a href="https://journals.aps.org/pra/abstract/10.1103/PhysRevA.110.022403" rel="noopener noreferrer"&gt;Distributed quantum architecture search&lt;/a&gt; from Luzhou Li's group at Sun Yat-sen University.&lt;/p&gt;

&lt;p&gt;This class of work highlights the platform's infrastructure role. Researchers are not merely calling a fixed algorithm; they are building new search strategies, learning processes, and experimental protocols on top of TensorCircuit.&lt;/p&gt;

&lt;h3&gt;
  
  
  Quantum Chemistry and Fermionic Simulation
&lt;/h3&gt;

&lt;p&gt;The quantum chemistry ecosystem around TenCirChem further extends TensorCircuit's application boundary. Quantum chemistry and fermionic simulation typically require complex Hamiltonian construction, differentiable optimization, tensor-network representations, and high-performance simulation. They therefore provide a demanding test case for any scientific computing platform.&lt;/p&gt;

&lt;p&gt;Representative works include &lt;a href="https://journals.aps.org/prresearch/abstract/10.1103/PhysRevResearch.5.023046" rel="noopener noreferrer"&gt;Efficient quantum simulation of electron-phonon systems by variational basis state encoder&lt;/a&gt; from teams at Tsinghua University and The Chinese University of Hong Kong, Shenzhen, as well as &lt;a href="https://pubs.acs.org/doi/abs/10.1021/acs.jctc.4c00200" rel="noopener noreferrer"&gt;Fast Emulation of Fermionic Circuits with Matrix Product States&lt;/a&gt; from Garnet Chan's group at Caltech. These studies show that the TensorCircuit ecosystem has moved from general quantum circuit simulation into more specialized domains such as quantum chemistry.&lt;/p&gt;

&lt;h3&gt;
  
  
  Quantum Sensing and Imaging
&lt;/h3&gt;

&lt;p&gt;TensorCircuit has also been used in quantum sensing, imaging, and experiment-facing tasks. Examples include &lt;a href="https://www.nature.com/articles/s41534-024-00914-w" rel="noopener noreferrer"&gt;End-to-end variational quantum sensing&lt;/a&gt; from Roger Melko's group at the Perimeter Institute, and &lt;a href="https://www.nature.com/articles/s42005-023-01290-1" rel="noopener noreferrer"&gt;Practical advantage of quantum machine learning in ghost imaging&lt;/a&gt; from Guihua Zeng's group at Shanghai Jiao Tong University. These works illustrate the platform's potential in quantum sensing and measurement-related applications.&lt;/p&gt;

&lt;p&gt;The value of a research platform is not captured by a single paper. It is reflected in its ability to support many research directions over time. More than 170 citing works, users across high-level institutions, and multiple examples in leading journals and conferences form an evidence chain that is stronger than any single promotional claim.&lt;/p&gt;

&lt;h2&gt;
  
  
  4. Has It Supported Industry-Oriented Applications?
&lt;/h2&gt;

&lt;p&gt;Academic citations show whether a platform can support research. Industry-oriented application cases show whether it can move toward real-world problems.&lt;/p&gt;

&lt;p&gt;It is important to be precise here. Quantum computing is still exploratory in many industrial contexts, so the right question is not whether it has already replaced classical solutions at scale. The better question is whether researchers and engineering teams in different fields have used the platform to build prototypes, workflows, and validation pipelines for real problem domains. From this perspective, TensorCircuit's application spillover already reaches multiple sectors.&lt;/p&gt;

&lt;p&gt;In agricultural diagnostics, researchers have used a quantum vision transformer for tomato leaf disease detection; see &lt;a href="https://eej.aut.ac.ir/article_5597.html" rel="noopener noreferrer"&gt;Enhancing Agricultural Diagnostics: Tomato Leaf Disease Detection Using Quantum Vision Transformer&lt;/a&gt;. In neuroscience and medical imaging, related works include &lt;a href="https://www.mdpi.com/2076-3425/14/4/401" rel="noopener noreferrer"&gt;Predicting Brain Age and Gender from Brain Volume Data Using Variational Quantum Circuits&lt;/a&gt; and &lt;a href="https://ieeexplore.ieee.org/abstract/document/10821329" rel="noopener noreferrer"&gt;Expanding the Horizon: Enabling Hybrid Quantum Transfer Learning for Long-Tailed Chest X-Ray Classification&lt;/a&gt;. In drug discovery, &lt;a href="https://www.nature.com/articles/s41598-024-67897-8" rel="noopener noreferrer"&gt;A hybrid quantum computing pipeline for real world drug discovery&lt;/a&gt; explores a hybrid quantum computing workflow for real drug discovery problems.&lt;/p&gt;

&lt;p&gt;TensorCircuit-NG has also appeared in security, communications, optimization, and computing systems. In software security, researchers have proposed lightweight quantum convolutional neural networks for malicious code detection. In drone and radar applications, hybrid quantum neural networks have been explored for radar return signal processing. In edge computing, quantum reinforcement learning has been used for joint resource allocation and task offloading. In finance, improved QAOA methods based on conditional value-at-risk have been studied for portfolio optimization. The significance of these cases is that they move quantum software frameworks from "quantum algorithm papers" into concrete domains such as agriculture, medicine, security, communications, finance, and drug discovery.&lt;/p&gt;

&lt;p&gt;External recognition provides additional context for the ecosystem. TensorCircuit has appeared in PhotonBox's 2022 list of influential quantum industry events in China, was listed as a recommended quantum software project in Google Summer of Code 2023, was used by NVIDIA in cuQuantum evaluation materials, was invited to participate in UnitaryHack 2024, and participated in Open Source Promotion Plan 2025. These forms of recognition do not replace technical validation, but they do show that TensorCircuit is not an isolated lab project. It has entered the public view of the open-source quantum software and high-performance computing ecosystems.&lt;/p&gt;

&lt;p&gt;Industrial maturity does not happen overnight. It typically moves from research prototypes, to open tools, to cross-domain collaboration, to engineering validation, and eventually to deployment. TensorCircuit-NG's current value lies in providing a reusable low-level toolchain for that process.&lt;/p&gt;

&lt;h2&gt;
  
  
  5. Does It Continue to Evolve?
&lt;/h2&gt;

&lt;p&gt;One defining feature of scientific infrastructure is that it is never finished.&lt;/p&gt;

&lt;p&gt;New hardware appears. New algorithms appear. New scientific demands appear. This makes sustained iteration more important than a single innovation.&lt;/p&gt;

&lt;p&gt;TensorCircuit's history is a good example. The project was first released in April 2020. From 2020 to 2021, TensorCircuit completed its core architecture, automatic differentiation mechanism, and early quantum algorithm modules, establishing the academic foundation for a unified tensor-computing framework. From 2021 to 2024, under the Apache License 2.0, the project continued to evolve in engineering: performance optimization, interface standardization, multi-backend support, and community ecosystem development gradually turned it into an open-source platform for global research users and developers.&lt;/p&gt;

&lt;p&gt;Since the launch of TensorCircuit-NG, or "Next Generation", in 2024, the project has moved beyond a quantum computing software framework toward a broader next-generation research infrastructure. It explores deeper integration among quantum computing, supercomputing, and intelligent computing, while continuing to expand its ecosystem in AI4Science and related areas.&lt;/p&gt;

&lt;p&gt;Sustained iteration is also visible in upstream and downstream ecosystem contributions. Upstream, core developers have contributed to standard machine learning frameworks such as TensorFlow, including work related to the automatic differentiation formula for complex-valued singular value decomposition and fixes to vectorized matrix multiplication. In the tensor-network ecosystem, TensorNetwork-NG continues to maintain the original Google TensorNetwork framework and keep it usable. Downstream, TenCirChem extends TensorCircuit capabilities into quantum computational chemistry workflows.&lt;/p&gt;

&lt;p&gt;These upstream and downstream contributions show that TensorCircuit-NG does not confine itself to a single framework. Instead, it builds connections among machine learning, tensor networks, quantum chemistry, and high-performance computing. This matters for Quantum x AI x HPC integration, because future research infrastructure cannot serve only one model family, one hardware type, or one class of algorithms.&lt;/p&gt;

&lt;p&gt;In the TC-NG architecture:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;quantum circuits;&lt;/li&gt;
&lt;li&gt;neural networks;&lt;/li&gt;
&lt;li&gt;tensor networks;&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;are brought into a unified computation-graph system.&lt;/p&gt;

&lt;p&gt;At the same time:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;CPUs;&lt;/li&gt;
&lt;li&gt;GPUs;&lt;/li&gt;
&lt;li&gt;HPC clusters;&lt;/li&gt;
&lt;li&gt;QPUs;&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;are becoming part of a unified resource pool.&lt;/p&gt;

&lt;p&gt;This marks a shift in platform positioning: from a quantum software framework to infrastructure for future scientific computing. Compared with projects that remain at the stage of concept demonstrations, short-term packaging, or slide-based roadmaps, more than six years of open-source development, continuous iteration, and repeated research-community validation say much more about a platform's real engineering capacity and long-term value.&lt;/p&gt;

&lt;h2&gt;
  
  
  Conclusion: What Builds Trust in Scientific Infrastructure?
&lt;/h2&gt;

&lt;p&gt;In the rapid development of Quantum x AI x HPC, industry narratives are converging.&lt;/p&gt;

&lt;p&gt;More and more platforms now talk about:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;AI4Science;&lt;/li&gt;
&lt;li&gt;hybrid quantum-classical computing;&lt;/li&gt;
&lt;li&gt;scientific research infrastructure.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;These directions are worth pursuing. But for users, researchers, and industry partners, the core criteria have not changed:&lt;/p&gt;

&lt;p&gt;Is the platform fully open source?&lt;/p&gt;

&lt;p&gt;Does it provide public benchmarks?&lt;/p&gt;

&lt;p&gt;Is it broadly and continuously used by high-quality research communities?&lt;/p&gt;

&lt;p&gt;Has it supported cross-industry application cases?&lt;/p&gt;

&lt;p&gt;Does it continue to evolve through sustained version updates?&lt;/p&gt;

&lt;p&gt;Once these questions are answered one by one, the value of a platform does not need to depend on slogans or conceptual messaging. For scientific infrastructure, long-term trust is built on verifiable code, reproducible experiments, growing academic adoption, application spillover into real problems, and engineering iteration that stands the test of time. In an era where technical narratives increasingly sound alike, these qualities are especially valuable.&lt;/p&gt;

</description>
      <category>ai</category>
      <category>quantum</category>
      <category>hpc</category>
    </item>
    <item>
      <title>TensorCircuit-NG: Quantum Software On AI, For AI, With AI</title>
      <dc:creator>Shixin Zhang</dc:creator>
      <pubDate>Wed, 27 May 2026 11:02:16 +0000</pubDate>
      <link>https://dev.to/refractionray/tensorcircuit-ng-quantum-software-on-ai-for-ai-with-ai-3mae</link>
      <guid>https://dev.to/refractionray/tensorcircuit-ng-quantum-software-on-ai-for-ai-with-ai-3mae</guid>
      <description>&lt;p&gt;Quantum computing and artificial intelligence are often discussed as two separate frontiers. One is about exploiting quantum mechanics for computation; the other is about building increasingly capable learning systems and agents. The core argument behind TensorCircuit-NG is that this separation is becoming less and less meaningful. If modern AI infrastructure has already solved core problems around automatic differentiation, compilation, accelerator execution, batching, and distributed training, then quantum software should stop reinventing those layers badly and start standing on top of them directly.&lt;/p&gt;

&lt;p&gt;This is the central idea behind TensorCircuit-NG. The project is a quantum software stack built in the age of AI, aimed at AI-facing workloads, and increasingly shaped for collaboration with AI agents. Its vision is simple: quantum software &lt;em&gt;on AI, for AI, with AI&lt;/em&gt;.&lt;/p&gt;

&lt;h2&gt;
  
  
  On AI: quantum software should inherit the AI stack
&lt;/h2&gt;

&lt;p&gt;Quantum software has long been held back by two familiar problems. Too much of the workload remains trapped in Python-level control flow or in classical state-vector simulation patterns that scale poorly. At the same time, many quantum libraries sit outside the deep learning ecosystems where most of the tooling innovation has happened. JAX, PyTorch, and TensorFlow already have mature answers to questions like compilation, vectorization, accelerator placement, and distributed execution, yet quantum software has often kept those capabilities at the edge of the stack.&lt;/p&gt;

&lt;p&gt;TensorCircuit-NG takes a different route. The framework treats quantum circuits as specialized tensor operations. That design choice opens up a large part of the AI toolchain almost “for free.” Automatic differentiation maps naturally onto variational quantum algorithms. Just-in-time compilation matters for repeated circuit evaluation. Vectorized mapping matters for batching over parameters, measurements, trajectories, or datasets. Accelerator support, mixed precision, and distributed execution are part of the design from the beginning.&lt;/p&gt;

&lt;p&gt;That philosophy shows up in the architecture. TensorCircuit-NG is built around a tensor-first worldview: every object is either a tensor or a network of tensors. Once that is the primitive, different computational models become easier to compose inside one workflow. Gate-based circuits, tensor networks, neural models, noisy simulators, analog evolution, approximate methods, and symbolic representations can live inside one coherent environment.&lt;/p&gt;

&lt;p&gt;The performance story follows directly from this design. TensorCircuit-NG supports both data parallelism and model parallelism across multiple devices and multiple hosts. In practice that means distribution over inputs, measurements, or noisy trajectories when the workload is embarrassingly parallel, and distribution over tensor-network slices when the contraction itself needs to be split across hardware. Benchmarks on both single-GPU and multi-GPU systems show that high-level Python APIs can still deliver high performance when the compilation and tensor-network substrate are done well.In representative workloads, that performance has reached speedups of several orders of magnitude over mainstream stacks such as IBM's Qiskit and Google's TensorFlow Quantum. &lt;/p&gt;

&lt;p&gt;TensorCircuit-NG acts as a bridge among quantum computing, high-performance computing, and intelligent computing. It also serves as an interface layer where quantum models can coexist with the rest of modern computational science. Researchers who want to embed quantum layers inside larger machine learning systems should be able to do so inside the same workflow, without crossing ecosystem boundaries every time the problem gets interesting.&lt;/p&gt;

&lt;h2&gt;
  
  
  For AI: a platform for fast quantum machine learning
&lt;/h2&gt;

&lt;p&gt;This is where the infrastructure becomes immediately useful. Quantum machine learning sits right at the intersection of circuit design, optimization, data pipelines, and repeated simulation. It is a workload that punishes slow software. If researchers want to try new ansatzes, change encodings, run ablations, train over many seeds, or sweep hyperparameters, then fast prototyping and efficient simulation matter more than slogans about QML.&lt;/p&gt;

&lt;p&gt;TensorCircuit-NG provides a strong platform for exactly this kind of work. Differentiable circuits, JIT compilation, batching, accelerator support, and distributed execution all live inside one environment. That makes it much easier to move from an idea for a QML model to a runnable prototype, and from a prototype to a meaningful simulation campaign.&lt;/p&gt;

&lt;p&gt;The scientific motivation for QML also becomes clearer in this setting. Attention shifts away from isolated benchmark wins and toward how quantum models behave on problems that already hurt classical AI. In our own work, this has already led to two systematic studies: one on bad data, and one on changing data.&lt;/p&gt;

&lt;p&gt;The first studies robustness. When labels are noisy, data is poisoned, or part of the training set later needs to be removed, quantum models may show a more favorable degradation profile and may be easier to unlearn. The second studies plasticity. In continual-learning settings, quantum models may preserve the ability to absorb new tasks for longer instead of becoming rigid.&lt;/p&gt;

&lt;p&gt;These are still open research questions. For a software project, though, the main point is straightforward: if people want to explore QML seriously, they need a platform that makes rapid iteration cheap. TensorCircuit-NG is meant to be that platform. It gives researchers a practical environment for fast QML prototyping, efficient simulation, and large-scale testing of ideas about robustness, unlearning, and adaptation.&lt;/p&gt;

&lt;h2&gt;
  
  
  With AI: a platform for agent-driven research
&lt;/h2&gt;

&lt;p&gt;The same logic carries over to AI agents. Once a scientific software stack is fast, structured, and composable, it becomes a natural substrate for agent-driven development. Agents are useful only when they can read real code, run real tools, inspect results, and keep iterating inside a live repository. That makes software design itself part of the agent story.&lt;/p&gt;

&lt;p&gt;TensorCircuit-NG is built with that use case in mind. The APIs are relatively concise, the examples and tests provide dense reference material, and the repository includes explicit rules and task-specific workflows for AI assistants. This lowers the cost of turning natural-language intent into runnable code, benchmarks, figures, and documentation.&lt;/p&gt;

&lt;p&gt;The project also ships built-in skills that push this further:&lt;/p&gt;

&lt;ul&gt;
&lt;li&gt;
&lt;code&gt;arxiv-reproduce&lt;/code&gt;, which turns a paper identifier into a reproduction workflow;&lt;/li&gt;
&lt;li&gt;
&lt;code&gt;performance-optimize&lt;/code&gt;, which injects optimization patterns such as &lt;code&gt;scan&lt;/code&gt;, &lt;code&gt;jit&lt;/code&gt;, &lt;code&gt;vmap&lt;/code&gt;, and contraction tuning;&lt;/li&gt;
&lt;li&gt;
&lt;code&gt;tc-rosetta&lt;/code&gt;, which translates code from other quantum frameworks with attention to intent rather than syntax alone;&lt;/li&gt;
&lt;li&gt;
&lt;code&gt;tutorial-crafter&lt;/code&gt;, which converts programs into polished narrative tutorials.&lt;/li&gt;
&lt;li&gt;and many more.&lt;/li&gt;
&lt;/ul&gt;

&lt;p&gt;Taken together, these tools make the framework a software platform where researchers can move from idea to prototype, from prototype to benchmark, and from benchmark to documentation with much less friction. That is the practical meaning of “with AI” here: TensorCircuit-NG is designed to work well with agents as a real development interface, not just as a chatbot wrapped around the codebase.&lt;/p&gt;

&lt;h2&gt;
  
  
  The deeper claim
&lt;/h2&gt;

&lt;p&gt;Taken together, these ideas add up to a stack-level thesis about the future of computational research.&lt;/p&gt;

&lt;p&gt;First, quantum software should no longer be architected as an isolated niche. It should inherit the best ideas from the AI and HPC worlds and expose them through abstractions that remain mathematically faithful to quantum workloads.&lt;/p&gt;

&lt;p&gt;Second, that same software stack should provide a strong platform for fast QML prototyping and efficient simulation, so ideas about robustness, unlearning, and continual adaptation can be tested quickly at scale.&lt;/p&gt;

&lt;p&gt;Third, the arrival of capable software agents changes the design target for scientific frameworks. A good framework now has to work well for skilled humans and also be understandable, navigable, and productively extensible for agents operating over the entire repository and toolchain.&lt;/p&gt;

&lt;p&gt;This is how TensorCircuit-NG understands itself: quantum software on AI, for AI, and with AI. It is built on the modern AI execution model, aimed at AI-relevant scientific questions, and increasingly shaped to participate in agent-mediated research workflows.&lt;/p&gt;

&lt;h2&gt;
  
  
  Getting started
&lt;/h2&gt;



&lt;div class="highlight js-code-highlight"&gt;
&lt;pre class="highlight shell"&gt;&lt;code&gt;pip &lt;span class="nb"&gt;install &lt;/span&gt;tensorcircuit-ng
&lt;/code&gt;&lt;/pre&gt;

&lt;/div&gt;



&lt;p&gt;An agent-first workflow also works well: ask your coding agent to install &lt;code&gt;tensorcircuit-ng&lt;/code&gt; and start building a small quantum application from natural-language instructions.&lt;/p&gt;

</description>
      <category>ai</category>
      <category>quantum</category>
    </item>
  </channel>
</rss>
