jax
40836799 - [Pallas:SC] Add new tpu_sc primitive for cumulative sum.

Commit

147 days ago

[Pallas:SC] Add new tpu_sc primitive for cumulative sum. Binds an SC implementation for lax.cumsum_p. PiperOrigin-RevId: 810788238

Author

brianwa84

Committer

Google-ML-Automation

Parents