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 _ListWithW since it is not needed anymore

Browse Source 
PiperOrigin-RevId: 510495372
This commit is contained in:
Yash Katariya 2023-02-17 12:29:29 -08:00 committed by jax authors
parent 031d15ed2d
commit 7a09fb98ef
1 changed files with 1 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
jax/_src/pjit.py
View File

@ -735,9 +735,6 @@ def pjit(
has_explicit_sharding)
class _ListWithW(list):
__slots__ = ('__weakref__',)
def hashable_pytree(pytree):
vals, treedef = tree_flatten(pytree)
vals = tuple(vals)
@ -941,7 +938,7 @@ def _create_pjit_jaxpr(fun, global_in_avals, api_name):
else:
jaxpr = core.ClosedJaxpr(jaxpr, consts)
final_consts = []
return _ListWithW([jaxpr, final_consts, global_out_avals])
return jaxpr, final_consts, global_out_avals
@lru_cache(maxsize=4096)