flax

[NVIDIA] Extend the custom fp8 accumulate dtype in non-jit scenarios

#3827

Merged