jax

57a4f6a0 - Update transpose rules of all primitives that define an unreduced_rule to call `.to_cotangent_aval()` on residuals.