We read every piece of feedback, and take your input very seriously.
To see all available qualifiers, see our documentation.
There was an error while loading. Please reload this page.
1 parent 9a686e0 commit be5d13aCopy full SHA for be5d13a
jax/_src/util.py
@@ -668,17 +668,12 @@ def wrapper(cls):
668
669
exclude_methods = {'__module__', '__dict__', '__doc__'}
670
671
- originals = {}
672
for attr_name, attr in cls.__dict__.items():
673
if attr_name not in exclude_methods:
674
- if hasattr(_original_func(attr), "_use_cpp"):
675
- originals[attr_name] = attr
676
- else:
+ if not hasattr(_original_func(attr), "_use_cpp"):
677
setattr(cpp_cls, attr_name, attr)
678
679
cpp_cls.__doc__ = cls.__doc__
680
- # TODO(pschuh): Remove once fastpath is gone.
681
- cpp_cls._original_py_fns = originals
682
return cpp_cls
683
684
return wrapper
0 commit comments