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

Merged

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

copybara-service merged 2 commits into google:main from kaixih:fp8_meta_custom_dtype_non_jit

kaixih force pushed 2 years ago

Update the type rule to support non-jit use cases

fd9c467f

kaixih force pushed to fd9c467f 2 years ago

IvyZX assigned

IvyZX 2 years ago

Work with older jax versions

5dfc0dbc

IvyZX approved these changes on 2024-04-23

IvyZX added pull ready

copybara-service merged 78189324 into main 2 years ago

Reviewers

IvyZX

Assignees

IvyZX

Labels

pull ready

Milestone

No milestone