$ the-wire · showcase
JAX hijax flags read at call time, Spanner stops deleting multiplexed sessions
By RepoJournal · Filed · About Google · Composed from the cited sources · methodology
JAX fixed a config-timing bug that silently disabled the custom_jvp3/custom_vjp3/remat3 experimental paths for import-time decorators, while Cloud Spanner clients dropped a session delete that the server rejects outright.
[hijax] read the custom_jvp3, custom_vjp3 and remat3 flags at call time google/jax
jax.custom_jvp, jax.custom_vjp and jax.checkpoint used to pick their classic or hijax implementation when the decorator ran, so functions decorated at import time (jax.nn.relu and dozens of other library functions) ignored jax.config.update and jtu.with_config unless JAX_CUSTOM_JVP3 was set before importing jax. They now read the flags at call time, which also means CustomJVP3Test and friends f...
[custom_jvp3] fix in-tree test failures under jax_custom_jvp3 google/jax
With JAX_CUSTOM_JVP3=1, 98 in-tree tests failed; that count is down to one, custom_api_test's test_def_method_forwarding_all_permutations, which the call-time flag read fixes. The largest failing group came from hijax primitives nested inside reduction bodies, such as cumlogsumexp lowering to a reduce_window whose body is logaddexp, and the default path is unchanged.
fix(spanner): stop deleting multiplexed sessions on rotation and close (#18332) googleapis/google-cloud-python
Multiplexed Cloud Spanner sessions cannot be deleted by clients, and calling DeleteSession on one returns INVALID_ARGUMENT, so old_session.delete() is gone from _rotate_multiplexed_session() and session_to_delete.delete() is gone from close() in both database_sessions_manager.py and its _async counterpart, which now simply clears the session reference.
perf(spanner): support PartialResultSet.last with background stream draining googleapis/google-cloud-python
Spanner marks the final PartialResultSet chunk with last = True, and the client previously blocked synchronously waiting for gRPC trailers and EOF frames before handing back the final rows. The client now returns as soon as it sees last = True and drains trailing metadata in the background, which also avoids streams abandoned early being marked CANCELLED by gRPC's C-core finalizer at garbage co...
Support generalized Hermitian eigenproblems in scipy.linalg.eigh google/jax
jax.scipy.linalg.eigh now handles type-1 generalized Hermitian eigenproblems via Cholesky reduction, with broadcast-compatible batch dimensions, JIT and vmap support, and differentiable eigenvalues and eigenvectors. Both inputs are symmetrized consistently with the standard eigh path, the eigenvalues-only path skips the eigenvector back-transform, and the b=None path is unchanged. Elsewhere: py...