Skip to content

Commit

Permalink
Add segment_ids support to pallas flash attention on TPU.
Browse files Browse the repository at this point in the history 
PiperOrigin-RevId: 572172125
  • Loading branch information
jax authors committed Oct 10, 2023
1 parent f2cda73 commit 60b77bd
Showing 1 changed file with 382 additions and 47 deletions.

0 comments on commit 60b77bd

Please sign in to comment.