nudge_diag

nudge_diag(mat, eps=1e-12)[source]

Stabilize matrix inversion by adding a tiny epsilon to the diagonal. Fully stable and smoothly differentiable in JAX. Supports unbatched (2D) and batched (N-D) square matrices.

Parameters:

mat (Array)
eps (float)

Return type:

Array