jax
02c19e96 - Make `jax.grad` and `compute_on` work correctly. If the forward pass has annotation to execute on CPU, then it's backward pass also executes on CPU.

Commit

2 years ago

Make `jax.grad` and `compute_on` work correctly. If the forward pass has annotation to execute on CPU, then it's backward pass also executes on CPU. PiperOrigin-RevId: 634917402

References

#21257 - Make `jax.grad` and `compute_on` work correctly. If the forward pass has annotation to execute on CPU, then it's backward pass also executes on CPU.

Author

yashk2810

Committer

a-googler

Parents

1043e24f

jax 02c19e96 - Make `jax.grad` and `compute_on` work correctly. If the forward pass has annotation to execute on CPU, then it's backward pass also executes on CPU.

jax
02c19e96 - Make `jax.grad` and `compute_on` work correctly. If the forward pass has annotation to execute on CPU, then it's backward pass also executes on CPU.