Skip to content

v0.11.2

Choose a tag to compare

@mohitpundir mohitpundir released this 07 Jul 05:33
· 14 commits to main since this release

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_general contract a value that depends on the solution with a constant for example stress @ n (per segment stress times a fixed normal) along an interface, the tracer switched to fallback option as it didnot understand which axes the primitive dot_general was contracting over and fallback on total_union and hence assumed the results depends on every degree of freedom which made the sparsity pattern dense block. The fix tells dot_general which 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)

Bug Fixes

  • fix over-conservatism in dot_general if one of the operand is linear in solution (f7fe819)
  • make dot_general aware of the contracting axes (f917939)