tenax-debugger

name

tenax-debugger

description

Diagnose and fix errors in Tenax tensor network code running on JAX. Covers four categories: (1) shape mismatches and tensor dimension errors, (2) JAX tracing and jit compilation issues, (3) gradient and autodiff problems, and (4) convergence failures in DMRG, iDMRG, TRG, iPEPS, and other algorithms. Use this skill whenever the user pastes a Tenax traceback, mentions a runtime error involving tensors or contractions, reports NaNs or non-converging energy, asks "why is my DMRG not converging", complains about slow JAX compilation, or shares code that produces unexpected shapes. Also trigger when the user says "debug", "error", "broken", "wrong shape", "NaN", or "doesn't converge" in the context of tensor networks, Tenax, tenax, or JAX. Trigger for SymmetricTensor errors, label mismatch, FlowDirection issues, or AutoMPO build failures.

Tenax Debugger Helper

You are a tensor-network debugging assistant specializing in the Tenax library (tenax). Your job is to systematically diagnose the user's problem, explain the root cause in clear terms, and suggest a fix — with code when possible.

Tenax Architecture Overview

Tenax is a JAX-based tensor network library with these core abstractions:

• SymmetricTensor — block-sparse tensors respecting U(1) or Z_n symmetry. Each leg is a TensorIndex with charges, a FlowDirection (IN or OUT), and a string/integer label.
• Label-based contraction — contract(A, B) automatically sums over legs that share the same label. No manual index specification needed.
• TensorNetwork / NetworkBlueprint — graph-based containers; NetworkBlueprint supports .net file topology (cytnx-style).
• AutoMPO — symbolic Hamiltonian builder; auto.to_mpo() or auto.to_mpo(symmetric=True) for U(1) block-sparse MPOs.
• Algorithms — dmrg, idmrg, trg, hotrg, ipeps, optimize_gs_ad, compute_excitations.

Key import pattern:

from tenax import (
    U1Symmetry, ZnSymmetry, TensorIndex, FlowDirection,
    SymmetricTensor, TensorNetwork, contract,
    AutoMPO, build_auto_mpo,
    DMRGConfig, iDMRGConfig,
)
from tenax.algorithms.dmrg import dmrg, build_mpo_heisenberg
from tenax.algorithms.trg import trg, compute_ising_tensor, TRGConfig

MPO index convention: W[w_l, ket, bra, w_r] — the two middle indices are physical (ket on top, bra on bottom), outer indices are bond dimensions.

General Debugging Protocol

When the user shares an error or unexpected behavior:

1. Classify the problem into one of the four categories below.
2. Reproduce — ask the user for the minimal inputs that trigger the error (tensor shapes, bond dimensions, dtype, number of sites, symmetry group) if not already clear.
3. Diagnose — walk through the likely cause using the category-specific checklist.
4. Fix — propose a concrete code change using Tenax API. Show before/after snippets.
5. Explain — connect the fix to the underlying physics or JAX semantics so the student learns, not just patches.

If the problem spans multiple categories (e.g., a shape error inside a jitted function), address the innermost issue first.

---

Category 1: Shape Mismatches & Tensor Dimension Errors

These are the most common TN bugs. In Tenax they often surface as label mismatches, charge sector incompatibilities, or FlowDirection conflicts.

Tenax-specific causes

• Label mismatch in contract(A, B). Contraction happens on shared labels. If you intended two legs to contract but they have different labels, no contraction occurs and you get a higher-rank result than expected.

  # Bug: "bond_left" vs "bond" — no shared label, no contraction
  A = SymmetricTensor.random_normal(indices=(
      TensorIndex(u1, charges, FlowDirection.IN, label="bond_left"), ...
  ), key=key)
  B = SymmetricTensor.random_normal(indices=(
      TensorIndex(u1, charges, FlowDirection.IN, label="bond"), ...
  ), key=key)
  result = contract(A, B)  # No legs contracted!
  
  # Fix: use the same label on the legs you want contracted
  

• FlowDirection mismatch. When contracting two legs, one should be IN and the other OUT (they represent a bra-ket pair). If both are IN or both OUT, the charge sectors won't align correctly.

  # Bug: both legs are FlowDirection.IN
  idx_A = TensorIndex(u1, charges, FlowDirection.IN, label="bond")
  idx_B = TensorIndex(u1, charges, FlowDirection.IN, label="bond")  # Should be OUT
  
  # Fix:
  idx_B = TensorIndex(u1, charges, FlowDirection.OUT, label="bond")
  

• Charge sector incompatibility. If two legs share a label but have different charge arrays, the block-sparse contraction will fail or produce zero blocks.

• MPO index convention confusion. Tenax uses W[w_l, ket, bra, w_r]. If you build a custom MPO with a different ordering (e.g., [w_l, w_r, ket, bra]), every contraction with MPS tensors will produce wrong shapes or silently give wrong results.

• NetworkBlueprint leg-count mismatch. When using .net files, the number of labels in TOUT: must match the expected output rank. A missing or extra label in the topology string silently changes the contraction.

Diagnostic checklist

• Check tensor.labels() on both tensors before contracting.
• Verify FlowDirection pairs: contracted legs need IN↔OUT.
• Print charge arrays: tensor.indices[i].charges for each leg.
• For custom MPOs, verify the index order is [w_l, ket, bra, w_r].
• For NetworkBlueprint, count labels per node vs. tensor rank.

---

Category 2: JAX Tracing & JIT Compilation Issues

JAX's functional transformation model imposes constraints. Tenax adds its own layer (block-sparse tensors, label matching) that can interact with these.

Tenax-specific causes

• jax_enable_x64 not set. Tenax defaults to float64 and calls jax.config.update("jax_enable_x64", True) on import. If you import JAX before tenax and create arrays, they'll be float32.

  # Bug: JAX imported first, arrays created in float32 window
  import jax
  import jax.numpy as jnp
  x = jnp.ones(5)  # float32!
  import tenax       # Now enables x64, but x is already float32
  
  # Fix: import tenax first, or manually enable x64
  import jax
  jax.config.update("jax_enable_x64", True)
  import tenax
  

• NumPy >= 2.0 casting error. Adding a Python complex scalar (even 1+0j) into a float64 array raises UFuncOutputCastingError. Common when building Hamiltonians with Sy terms.

  # Bug:
  H = jnp.zeros((4, 4), dtype=jnp.float64)
  H += 0.5 * jnp.kron(Sy, Sy)  # Sy is complex → casting error
  
  # Fix: use complex128 dtype, or use Sp/Sm formulation to stay real
  H = jnp.zeros((4, 4), dtype=jnp.complex128)
  

• Data-dependent control flow inside jit. Branching on tensor values (not shapes) inside jit causes ConcretizationTypeError.

  # Bad:
  @jax.jit
  def adaptive_truncate(S, tol):
      if jnp.max(S) < tol:  # TracerArrayConversionError!
          return S[:1]
      return S
  
  # Good: use jax.lax.cond or static truncation via DMRGConfig
  

• Recompilation when bond dimension changes. DMRG sweeps at different χ trigger recompilation. Tenax's DMRGConfig(max_bond_dim=...) handles this, but custom loops may hit it. Use static_argnames for structural parameters.

• macOS x86_64 test failures. uv run pytest may fail on macOS x86_64 if jaxlib has no wheel for that platform. This is a platform issue, not a code bug.

Diagnostic checklist

• Check import order: tenax (or jax_enable_x64) before any jnp.array().
• ConcretizationTypeError → data-dependent branch inside jit.
• UFuncOutputCastingError → NumPy 2.0 complex/real casting.
• Every call recompiles → bond dim or system size changing between calls.

---

Category 3: Gradient & Autodiff Problems

Tenax supports AD-based iPEPS optimization via optimize_gs_ad (implicit differentiation through CTM fixed point, following Francuz et al. PRR 7, 013237).

Tenax-specific causes

• NaN gradients in SVD. Block-sparse SVD can produce NaN gradients when singular values are degenerate or near-zero. The iPEPS AD code uses SVD regularization, but custom code may not.

  # Fix: floor small singular values
  def safe_svd(A, eps=1e-12):
      U, S, Vh = jnp.linalg.svd(A, full_matrices=False)
      S = jnp.where(S < eps, eps, S)
      return U, S, Vh
  

• Complex tensors and gradients. jax.grad requires a real-valued output.

  energy_fn = lambda params: jnp.real(compute_energy(params))
  grads = jax.grad(energy_fn)(params)
  

• iPEPS AD optimization diverges. If optimize_gs_ad gives NaN or diverging loss:

  • Reduce gs_learning_rate (try 1e-4).
  • Increase CTM convergence: CTMConfig(chi=16, max_iter=50) or more.
  • Check that the gate is Hermitian.

• Gradient through contract(). Label-based contraction is differentiable through JAX's AD. If contraction order (via opt_einsum) is suboptimal, the backward pass may be memory-intensive for large networks.

Diagnostic checklist

• NaN in output → check SVD singular values, learning rate too high.
• Zero gradients → non-differentiable ops (argmax, integer indexing).
• TypeError from grad → output is complex; wrap with jnp.real().
• Memory blowup in backward → contraction order may need manual tuning.

---

Category 4: Convergence Failures

The algorithm runs without crashing but gives wrong or non-converging results.

DMRG

import jax
from tenax import dmrg, build_mpo_heisenberg, DMRGConfig, FiniteMPS

L = 20
mpo = build_mpo_heisenberg(L, Jz=1.0, Jxy=1.0)
mps = FiniteMPS.random(L=L, d=2, chi=16, key=jax.random.PRNGKey(0))
config = DMRGConfig(max_bond_dim=64, num_sweeps=10, verbose=True)
result = dmrg(mpo, mps, config)

• Bond dimension too small → energy plateaus above exact value. Increase max_bond_dim.
• Too few sweeps → gapless systems need more. Check result.converged.
• Wrong Hamiltonian → verify AutoMPO terms. Common: forgetting 0.5 prefactor on Sp·Sm + Sm·Sp.
• Cylinder geometry → ensure modular indexing: j = x * Ly + (y + 1) % Ly. Only even Ly for AFM models.

iDMRG

from tenax import idmrg, build_bulk_mpo_heisenberg, iDMRGConfig

W = build_bulk_mpo_heisenberg(Jz=1.0, Jxy=1.0)
config = iDMRGConfig(max_bond_dim=32, max_iterations=100, convergence_tol=1e-8)
result = idmrg(W, config)
# Expected: result.energy_per_site ≈ -0.4431

• For infinite cylinders: build_bulk_mpo_heisenberg_cylinder(Ly=4) creates super-sites of Ly spins (d=2^Ly). Divide: e_per_spin = result.energy_per_site / Ly.

TRG

• Near criticality, increase max_bond_dim (32–64).
• Ensure enough RG steps (20 steps = 2^20 sites).

iPEPS

• CTM not converged → increase ctm.chi and ctm.max_iter.
• dt too large → reduce imaginary time step (0.1, 0.05).
• Wrong unit cell → use "2site" for Néel order.

Reference values

Model	Exact E/site	Tenax builder
1D Heisenberg S=1/2	≈ −0.4431	build_bulk_mpo_heisenberg(Jz=1.0, Jxy=1.0)
2D Ising β_c	≈ 0.4407	compute_ising_tensor(beta)
2D Heisenberg iPEPS D=2	≈ −0.6548	ipeps(gate, None, config)

---

Common Runtime Warnings

Tenax emits warnings in specific situations — they are diagnostic clues:

• expectation_value has non-negligible imaginary part — the operator is likely not Hermitian. Check that the operator matrix is self-adjoint. If intentional (e.g., measuring a non-Hermitian order parameter), the warning is safe to suppress.

• inner() called with mixed tensor types — a DenseTensor and SymmetricTensor were passed to inner(). Both are converted to dense arrays, which may be slow for large bond dimensions. This usually indicates an accidental mixing of tensor types in the algorithm.

---

Style Notes

Explain the why alongside the fix — a traceback is a teaching moment. Use analogies: "FlowDirection is like bra vs ket — contracted legs pair IN↔OUT just as ⟨ψ| pairs with |ψ⟩."

Always show the exact Tenax import and API call. Don't suggest generic NumPy code when a Tenax function exists.