v0.11.2
This release adds handlers for missing primitives for JAX linear algebra operations at the element level, such as jnp.linalg.inv, jnp.linalg.eig and also for jnp.flip and jnp.stack.
This release also fixes the edge cases with dot_general such as:
- over-conservatism if one of the operands is just linearly dependent on the solution. Such a case was never checked for how the operand is dependent on the solution because only a nonlinear dependency will lead to a non-zero Hessian entry. The fix records an operand's self-second-order coupling only if it is nonlinear in the solution; the cross second-order coupling between operands is always kept.
- when we
dot_generalcontract a value that depends on the solution with a constant for examplestress @ n(per segment stress times a fixed normal) along an interface, the tracer switched to fallback option as it didnot understand which axes the primitivedot_generalwas contracting over and fallback ontotal_unionand hence assumed the results depends on every degree of freedom which made the sparsity pattern dense block. The fix tellsdot_generalwhich axis is being summed over, so dependencies are combined only along the axis and thus the pattern stays block-diagonal, and no real entry is lost.
This release also adds a warning for the primitives that are not considered.
Enhancements
- add handlers for eigvalues, flip (35c5652)
- add sparsity handlers for dense linear-algebra ops and stack (d1277a4)
- add warnings for not considered primitives (e519f3f)