jax
2d1479dc - Repurpose `tpu_sparsecore` annotation for offloading purposes

Commit

50 days ago

Repurpose `tpu_sparsecore` annotation for offloading purposes Today only gather/scatter are supported with this offload strategy. PiperOrigin-RevId: 826549779

References

#31381 - Remove old ROCm build code

#33008 - Repurpose `tpu_sparsecore` annotation for offloading purposes

#33157 - Resolve undefined behavior in bitshift unit test

#33186 - Make nvidia version data optional for ROCm builds

Author

vsytch

Committer

Google-ML-Automation

Parents

afbd5ff1

jax 2d1479dc - Repurpose `tpu_sparsecore` annotation for offloading purposes

jax
2d1479dc - Repurpose `tpu_sparsecore` annotation for offloading purposes