Recent versions of JAX (0.4.28+) have made some changes to:
Hashing of tracers;
Tree-map'ing over Nones;
Callbacks;
Pretty-printing.
With this update, we should now be compatible with both old and new versions of JAX: this fixes both some new crashes, and some new warnings. (#719, #724, #753, #758, thanks @jakevdp, @hawkinsp!)
Better errors
The error messages from eqx.error_if are now substantially more informative: they include traceback information including the stack, and mention the availability of the EQX_ON_ERROR variable. We also do a much better job hiding the large unhelpful printouts that XLA gives by default. (#785, #803)
The default value of EQX_ON_ERROR_BREAKPOINT_FRAMES is now 1. (#777) The impact of this is that using eqx.error_if alongside EQX_ON_ERROR=breakpoint will now:
reliably always open a debugger, rather than sometimes crashing at trace-time due to upstream JAX bug #16732.
however, by default the debugger will no longer include any additional stack frames above it (accessed via u).
much of the above is now explained in a printed-out informative message prior to the debugger opening.
Bugfixes
eqx.filter_{jacfwd, jacrev} now only apply filtering to their inputs but not their outputs. Previously this was problematic as there was no way to represent static-input-by-static-output in the returned Jacobian, so pieces were silently dropped. (#734, thanks @lockwo!)
eqx.tree_at can now be used to replace empty tuples. (#715, #717, #722, thanks @lockwo!)
eqx.filter_custom_jvp no longer raises a trace-time crash in some scenarios in which its **kwargs were erroneously counted as having tangents. (patrick-kidger/equinox#745, #749)
No longer getting a trace-time crash when doing a particular combination of vmap + autodiff + checkpointed while loops. This occurred when using optimistix.BFGS around diffrax.diffeqsolve. (#777)
Fixed a trace-time crash when:
using a checkpointed while loop...
...with a body function that has a closed-over tracer...
...and that closed-over tracer is differentiated...
...and there are no other closed-over tracers that are differentiated...
...and the dependency on that tracer is only linear.
Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting @dependabot rebase.
Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR:
- `@dependabot rebase` will rebase this PR
- `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it
- `@dependabot merge` will merge this PR after your CI passes on it
- `@dependabot squash and merge` will squash and merge this PR after your CI passes on it
- `@dependabot cancel merge` will cancel a previously requested merge and block automerging
- `@dependabot reopen` will reopen this PR if it is closed
- `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually
- `@dependabot show ignore conditions` will show all of the ignore conditions of the specified dependency
- `@dependabot ignore this major version` will close this PR and stop Dependabot creating any more for this major version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this minor version` will close this PR and stop Dependabot creating any more for this minor version (unless you reopen the PR or upgrade to it yourself)
- `@dependabot ignore this dependency` will close this PR and stop Dependabot creating any more for this dependency (unless you reopen the PR or upgrade to it yourself)
Bumps equinox from 0.11.4 to 0.11.5.
Release notes
Sourced from equinox's releases.
... (truncated)
Commits
ea1af82Version bump2188e01Runtime errors inside eqx.filter_jit are now very readable.68cc26aFix missingout_axesfor transpose-of-vprim.028a6f4Better strict error (#783)fe99597Equinox runtime errors now include a stack trace and more debug info.2d868a9Update contributing guide23b3042Fix for BFGS + diffeqsolve70d664cNow defaulting to providing only one frame to the on-error debugger.8efb0fcDoc fixesd3a01c8Fixed a case in which we're hashing tracers.Dependabot will resolve any conflicts with this PR as long as you don't alter it yourself. You can also trigger a rebase manually by commenting
@dependabot rebase.Dependabot commands and options
You can trigger Dependabot actions by commenting on this PR: - `@dependabot rebase` will rebase this PR - `@dependabot recreate` will recreate this PR, overwriting any edits that have been made to it - `@dependabot merge` will merge this PR after your CI passes on it - `@dependabot squash and merge` will squash and merge this PR after your CI passes on it - `@dependabot cancel merge` will cancel a previously requested merge and block automerging - `@dependabot reopen` will reopen this PR if it is closed - `@dependabot close` will close this PR and stop Dependabot recreating it. You can achieve the same result by closing it manually - `@dependabot show