jax

32e419d1 - Fix eigh JVP to ensure that both the primal and tangents of the eigen… (#3550)

Fix eigh JVP to ensure that both the primal and tangents of the eigen… (#3550) * Fix eigh JVP to ensure that both the primal and tangents of the eigenvalues are real. Add test to jax.test_util.check_jvp that ensure the primals and both the primals and tangents produced by a JVP rule have identical types. * Cast input to static indexing grad tests to a JAX array so new type check passes.