tenet.ad

Stabilized JAX differentiation rules for svd and eigh at a degenerate spectrum.

tenet.ad

Opt-in stabilized VJPs for svd/eigh on the JAX backend. Call install.

JAX's own SVD/eigh VJPs carry 1/(sigma_i - sigma_j) / 1/(w_i - w_j) factors that are NaN at exact degeneracy, and under a non-Abelian symmetry degeneracy inside a coupled sector is generic rather than an edge case. install() replaces them with the Lorentzian-broadened form, 1/x -> x/(x**2 + eps); uninstall() restores autoray's stock bindings.

Three things a caller must know:

• It is process-global for the JAX backend, by design. The seam is autoray.register_function("jax", "linalg.svd", ...), autoray's own extension point, so after install() any ar.do("linalg.svd", jax_array) in the process -- quimb's included -- gets the broadened VJP. Installation is therefore an explicit function call, not an import side effect like tenet.pytree's: mutating another library's dispatch table is the user's act, not an import's.

• The broadened gradient is correct, not merely finite, exactly when the objective is gauge-invariant on each degenerate subspace. At exact degeneracy the singular vectors within a multiplet are only defined up to a unitary, so dU/dA genuinely does not exist; what exists is the derivative of gauge-invariant combinations such as U S Vh, the singular values, or any projector onto the multiplet. Broadening returns the correct value for those and an eps-suppressed arbitrary value for the rest. It is a precondition, not a hidden approximation: a gauge-dependent objective gets an eps-dependent answer, which is to say a meaningless one.

• eps is in units of sigma squared -- safe_inverse(x) = x/(x**2 + eps) with x = sigma_j - sigma_i -- so a default tuned for an O(1)-normalized spectrum assumes a normalization our tensors do not guarantee. Hence the knob.

Broadening rather than a hard tolerance, because a hard cut makes the gradient discontinuous in the spectrum. qr, lq and polar need nothing of their own: their backend rules are already well defined where their inputs are full-rank.

tenet.enable_jax(ad=True) is the one-call spelling of the three statements below, and is what the docs teach. It is the same explicit act this docstring demands -- ad is opted into by name there, never defaulted on, precisely because of the process-global reach above -- and it runs this module's install() unchanged.

Usage::

import tenet

tenet.enable_jax(ad=True)  # pytree registration + the broadened VJPs

def objective(t):
    u, s, vh = tenet.linalg.svd(t)   # per coupled sector
    return tenet.norm(u @ s @ vh)

g = jax.grad(objective)(t)

install

install(*, epsilon: float = 1e-12) -> None

Swap JAX's stock linalg.svd/linalg.eigh VJPs for broadened ones.

Parameters:

Name	Type	Description	Default
epsilon	float	The Lorentzian broadening, in units of sigma squared; the default is the PRX value and assumes an O(1)-normalized spectrum. Read at trace time, so set it before the hot loop.	1e-12

Notes

Process-global for the JAX backend, by design -- see the module docstring. Idempotent.

uninstall

uninstall() -> None

Restore autoray's stock JAX bindings. Idempotent.

register_function writes _FUNCS[backend, name]; deleting that entry is what restores the stock binding, because the next lookup falls back through autoray's import machinery and reapplies autoray's own wrapper for the function. Registering a remembered "old" callable would instead leave a second override in place. Private attributes, knowingly: autoray ships an install seam and no matching uninstall one, and the only real caller is the test that needs stock and broadened gradients in the same process.