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

Convert tuple to DeviceAssignment on the replicated compilation path.

Browse Source 
PiperOrigin-RevId: 533258935
This commit is contained in:
jax authors 2023-05-18 14:52:37 -07:00
parent 08169291a4
commit 6034e87ddf
1 changed files with 0 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jax/_src/interpreters/pxla.py
View File

@ -2918,7 +2918,6 @@ def _compile_replicated_mesh_executable_from_hlo(
host_callbacks, has_unordered_effects, ordered_effects, kept_var_idx,
backend, da, committed, pmap_nreps, jaxpr_debug_info):
assert not auto_spmd_lowering
assert isinstance(da, _DeviceAssignment)
in_shardings = semantics_in_shardings.shardings
out_shardings = semantics_out_shardings.shardings