jax

1c32cfdc - Add a test for remat transpose when the primal inputs are reduced. Make sure the cotangents are unreduced