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

Fix ndarray comparision in sharding_impls.py

Browse Source
This commit is contained in:
Skye Wanderman-Milne 2023-04-12 20:43:57 +00:00
parent 49e68dbe80
commit e2f1e7d28e
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
jax/_src/sharding_impls.py
View File

@ -535,7 +535,7 @@ class PositionalSharding(XLACompatibleSharding):
return False
if id(self) == id(other):
return True
all_ids_equal = bool(np.all(self._ids == other._ids))
all_ids_equal = np.array_equal(self._ids,other._ids)
if id(self._devices) == id(other._devices) and all_ids_equal:
return True
return self._devices == other._devices and all_ids_equal