jax
60b77bdd - Add segment_ids support to pallas flash attention on TPU.

Commit

2 years ago

Add segment_ids support to pallas flash attention on TPU. PiperOrigin-RevId: 572172125

Author

a-googler

Committer

a-googler

Parents