set_jax_expm_max_squarings

static Config.set_jax_expm_max_squarings(max_squarings)

Set the maximum number of squarings used by the JAX matrix exponential.

Parameters:

max_squarings (int) – Positive integer limit for the JAX matrix exponential.

Note

Simphony’s JAX-based propagators ultimately call jax.scipy.linalg.expm(), which evaluates the matrix exponential using a scaling-and-squaring scheme together with a Padé approximation. The basic idea is to first reduce the matrix norm, compute an accurate approximation for the smaller matrix, and then undo the scaling.

In simplified form, the method uses

\[e^A = \left(e^{A / 2^s}\right)^{2^s},\]

where \(s\) is the number of squaring steps.

For the scaled matrix \(A / 2^s\), JAX approximates the exponential by a rational Padé form

\[e^{A / 2^s} \approx r_m(A / 2^s) = P_m(A / 2^s)\,Q_m(A / 2^s)^{-1},\]

with matrix polynomials \(P_m\) and \(Q_m\). In practice, the role of scaling is to make \(A / 2^s\) small enough that this rational approximation is accurate and numerically stable.

JAX estimates how much scaling is needed from the matrix norm and then performs the required number of repeated squarings. Roughly speaking,

\[s \approx \max\left(0, \left\lceil \log_2 \|A\|_1 \right\rceil - c\right),\]

where \(c\) is an internal dtype-dependent threshold. The max_squarings parameter is therefore not a direct accuracy knob for the Padé approximation itself; instead, it is an upper bound on how much scaling-and-squaring JAX is allowed to use.

If the estimated number of squarings exceeds this bound, JAX may refuse the requested evaluation and return nan values rather than silently producing an unreliable result. Increasing max_squarings makes larger-norm matrices admissible, but also permits more work in the exponential evaluation.

In Simphony this matters when the effective generator passed to expm becomes large, for example because of large time steps, strong couplings, large driving amplitudes, or Hamiltonians written in units that produce large matrix norms. If you encounter unexpected nan values in JAX-based propagation, this setting is one of the first things worth checking.