xla
[Pallas] Improve FlashAttention segment_ids test case
#7034

Merged

Commits

tmp

alanwaketan committed 1 year ago
some minor changes

alanwaketan committed 1 year ago
Fix everything

alanwaketan committed 1 year ago
Enable backward as well

alanwaketan committed 1 year ago
Delete customer code

alanwaketan committed 1 year ago
Improve test case

alanwaketan committed 1 year ago

Loading