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

Commit

343 days ago

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

References

#33008 - Repurpose `tpu_sparsecore` annotation for offloading purposes

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