New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Improve hashing situation, add xr.Dataset as pytree node, implement i… #4

Merged

allen-adastra merged 1 commit into main from allenw/ds_grad_more_tests

Sep 21, 2024

Owner

allen-adastra commented Sep 21, 2024

xr.Dataset is now a pytree node
Improve coordinate hashing logic
Make xr.IndexVariable static
Tests for a set of identity transforms
- Going back and forth between xr and the optional xj types
- flatten and unflatten
- jit on the identity function
- vmap on the identity function
- eqx.partition with all trues
- eqx.filter with all trues
Test for jax.grad on float Variable and DataArray types


          Improve hashing situation, add xr.Dataset as pytree node, implement i…

6272c2b

…dentity transform tests and grads

allen-adastra merged commit c5cdb33 into main

3 checks passed

allen-adastra deleted the allenw/ds_grad_more_tests branch

September 21, 2024 15:42

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet