Skip to content

Commit be5d13a

Browse files
hawkinspGoogle-ML-Automation
authored andcommitted
Remove code that preserved _original_py_fns on C++ classes.
This no longer appears to be used. PiperOrigin-RevId: 737715578
1 parent 9a686e0 commit be5d13a

File tree

1 file changed

+1
-6
lines changed

1 file changed

+1
-6
lines changed

jax/_src/util.py

Lines changed: 1 addition & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -668,17 +668,12 @@ def wrapper(cls):
668668

669669
exclude_methods = {'__module__', '__dict__', '__doc__'}
670670

671-
originals = {}
672671
for attr_name, attr in cls.__dict__.items():
673672
if attr_name not in exclude_methods:
674-
if hasattr(_original_func(attr), "_use_cpp"):
675-
originals[attr_name] = attr
676-
else:
673+
if not hasattr(_original_func(attr), "_use_cpp"):
677674
setattr(cpp_cls, attr_name, attr)
678675

679676
cpp_cls.__doc__ = cls.__doc__
680-
# TODO(pschuh): Remove once fastpath is gone.
681-
cpp_cls._original_py_fns = originals
682677
return cpp_cls
683678

684679
return wrapper

0 commit comments

Comments
 (0)