xla
ce1205e1 - [Pallas] Make repeat_with_fixed_output_size not OOM on VMEM (#7145)

Commit
1 year ago
[Pallas] Make repeat_with_fixed_output_size not OOM on VMEM (#7145) Summary: openxla.org/xla/operation_semantics#reducewindow doesn't support int64. Let's make sure input to cumsum is always int32. Test Plan: python test/test_gmm.py
Author
Parents
Loading