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

Allow unhashable callables in jax.eval_shape.

Browse Source 
PiperOrigin-RevId: 599691923
This commit is contained in:
Matthew Johnson 2024-01-18 19:14:38 -08:00 committed by jax authors
parent 6d21b498c0
commit 0d4f200e08
1 changed files with 3 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
jax/_src/api.py
View File

@ -2835,6 +2835,9 @@ def eval_shape(fun: Callable, *args, **kwargs):
>>> print(out.dtype)
float32
"""
# Workaround to support unhashable callables.
try: hash(fun)
except TypeError: fun = partial(fun)
# The traced_for name is `jit` so as to get maximum tracing cache hits.
# Eventually, we should deprecate `eval_shape` and expose it like AOT style.
f, dbg, res_paths, args_flat, _, out_tree, _, _ = pjit.get_wrapped_fun(