mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-19 05:16:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

Remove implicit sharding annotation for tpu custom call.

Browse Source 
PiperOrigin-RevId: 691876343
This commit is contained in:
jax authors 2024-10-31 11:29:27 -07:00
parent 8296f6e0ba
commit c758373b9c
1 changed files with 0 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/_src/tpu_custom_call.py
View File

@ -204,9 +204,6 @@ class CustomCallBackendConfig:
if i + 1 != len(self.flags):
config.write(b",")
config.write(b"]")
# Prevent the compiler from sharding the custom call beyond what Mosaic does
# based on user annotations
config.write(b', "implicit_sharding": {"type": "MANUAL"}')
config.write(b"}")
return config.getvalue()