Jacobians
At an equilibrium
\[
z^\star=f_\theta(z^\star,x),
\]
the local Jacobian with respect to the state is
\[
J_z = \frac{\partial f_\theta}{\partial z}(z^\star,x).
\]
This matrix controls stability, implicit gradients, and solver behavior.
The stochastic trace/norm diagnostic follows Hutchinson
[14], while its use as an
equilibrium regularizer connects to Jacobian-regularized DEQ
[6].
Full Jacobian
For small states, materialize
\[
J_z[i,j]
=
\frac{\partial f_i}{\partial z_j}.
\]
from silva_networks import full_jacobian
J = full_jacobian(lambda z: transition(z, x), z_star)
This is exact up to autograd precision, but memory scales like
\[
O(n^2)
\]
for a flattened state dimension \(n\).
Vector-Jacobian Product
For large states, a vector-Jacobian product computes
\[
J_z^\top v
\]
without constructing \(J_z\). This is the primitive used by implicit backward
passes and power iteration.
from silva_networks import vjp
jtv = vjp(lambda z: transition(z, x), z_star, v)
Jacobian-Vector Product
The companion product is
\[
J_z v.
\]
from silva_networks import jvp
value, jv = jvp(lambda z: transition(z, x), z_star, v)
Spectral Radius
Local contraction is governed by
\[
\rho(J_z)=\max_i |\lambda_i(J_z)|.
\]
When
\[
\rho(J_z)<1,
\]
the fixed point is locally stable for the linearized update. The package
estimates this quantity using VJP-based power iteration.
Hutchinson Norm Estimate
For a Rademacher vector \(v\), with entries sampled from \(\{-1,+1\}\),
\[
\mathbb E_v\|J_z^\top v\|_2^2=\|J_z\|_F^2.
\]
The helper hutchinson_jacobian_norm averages this quantity over several
probes.
StabilityReport
dataclass
Local fixed-point diagnostics evaluated at one state.
Source code in src/silva_networks/jacobian.py
| @dataclass(frozen=True)
class StabilityReport:
"""Local fixed-point diagnostics evaluated at one state."""
residual: float
spectral_radius: float
jacobian_norm_estimate: float
samples: int
|
full_jacobian
Materialize the full Jacobian for small states.
Source code in src/silva_networks/jacobian.py
| def full_jacobian(f: Callable[[Tensor], Tensor], z: Tensor) -> Tensor:
"""Materialize the full Jacobian for small states."""
_validate_state(z)
z_req = z.detach().requires_grad_(True)
def flat_f(flat_z: Tensor) -> Tensor:
return f(flat_z.reshape_as(z_req)).reshape(-1)
return torch.autograd.functional.jacobian(flat_f, z_req.reshape(-1), vectorize=False)
|
hutchinson_jacobian_norm
hutchinson_jacobian_norm(f, z, samples=8, squared=True)
Estimate ||J_f(z)||_F or its square with Rademacher VJP probes.
Source code in src/silva_networks/jacobian.py
| def hutchinson_jacobian_norm(
f: Callable[[Tensor], Tensor],
z: Tensor,
samples: int = 8,
squared: bool = True,
) -> Tensor:
"""Estimate ``||J_f(z)||_F`` or its square with Rademacher VJP probes."""
_validate_state(z)
if samples < 1:
raise ValueError("samples must be positive")
z_req = z.detach().requires_grad_(True)
y = f(z_req)
_validate_output(y, z)
if not y.requires_grad:
return torch.zeros((), device=z.device, dtype=z.dtype)
acc = torch.zeros((), device=z.device, dtype=z.dtype)
for _ in range(samples):
probe = torch.empty_like(y).bernoulli_(0.5).mul_(2.0).sub_(1.0)
(jtv,) = torch.autograd.grad(
y,
z_req,
probe,
retain_graph=True,
create_graph=True,
allow_unused=True,
)
if jtv is None:
continue
acc = acc + torch.sum(jtv * jtv)
estimate = acc / samples
return estimate if squared else torch.sqrt(estimate.clamp_min(0.0))
|
jvp
Return f(z) and J_f(z) v.
Source code in src/silva_networks/jacobian.py
| def jvp(f: Callable[[Tensor], Tensor], z: Tensor, v: Tensor) -> tuple[Tensor, Tensor]:
"""Return ``f(z)`` and ``J_f(z) v``."""
_validate_state(z)
if v.shape != z.shape:
raise ValueError("v must have the same shape as z for a JVP")
return torch.autograd.functional.jvp(f, z.detach(), v.detach(), create_graph=False)
|
spectral_radius
spectral_radius(step, z, iters=30)
Power iteration on J_step(z)^T using VJP calls.
Source code in src/silva_networks/jacobian.py
| def spectral_radius(step: Callable[[Tensor], Tensor], z: Tensor, iters: int = 30) -> float:
"""Power iteration on ``J_step(z)^T`` using VJP calls."""
_validate_state(z)
if iters < 1:
raise ValueError("iters must be positive")
z_req = z.detach().requires_grad_(True)
y = step(z_req)
_validate_output(y, z)
if not y.requires_grad:
return 0.0
v = torch.randn_like(z_req)
v = v / (torch.linalg.norm(v.reshape(-1)) + 1e-12)
rho = torch.zeros((), device=z.device, dtype=z.dtype)
for _ in range(iters):
(jtv,) = torch.autograd.grad(
y,
z_req,
v,
retain_graph=True,
create_graph=False,
allow_unused=True,
)
if jtv is None:
return 0.0
rho = torch.linalg.norm(jtv.reshape(-1))
v = jtv / (rho + 1e-12)
return float(rho.detach().cpu())
|
stability_report
stability_report(f, z, step=None, samples=8, iters=30)
Compute residual, spectral-radius, and Jacobian-norm diagnostics.
Source code in src/silva_networks/jacobian.py
| def stability_report(
f: Callable[[Tensor], Tensor],
z: Tensor,
step: Callable[[Tensor], Tensor] | None = None,
samples: int = 8,
iters: int = 30,
) -> StabilityReport:
"""Compute residual, spectral-radius, and Jacobian-norm diagnostics."""
if samples < 1:
raise ValueError("samples must be positive")
if iters < 1:
raise ValueError("iters must be positive")
with torch.no_grad():
residual = float(torch.linalg.norm((f(z) - z).reshape(-1)).detach().cpu())
step_fn = step or f
rho = spectral_radius(step_fn, z, iters=iters)
norm_est = float(hutchinson_jacobian_norm(f, z, samples=samples, squared=False).detach().cpu())
return StabilityReport(
residual=residual, spectral_radius=rho, jacobian_norm_estimate=norm_est, samples=samples
)
|
vjp
vjp(f, z, v, create_graph=False)
Return J_f(z)^T v without materializing the Jacobian.
Source code in src/silva_networks/jacobian.py
| def vjp(f: Callable[[Tensor], Tensor], z: Tensor, v: Tensor, create_graph: bool = False) -> Tensor:
"""Return ``J_f(z)^T v`` without materializing the Jacobian."""
_validate_state(z)
z_req = z.detach().requires_grad_(True)
y = f(z_req)
_validate_output(y, z, v)
if not y.requires_grad:
return torch.zeros_like(z_req)
(jtv,) = torch.autograd.grad(
y,
z_req,
v,
retain_graph=create_graph,
create_graph=create_graph,
allow_unused=True,
)
return torch.zeros_like(z_req) if jtv is None else jtv
|
Where to Go Next