Make eval_shape and invertible overridable

google · Feb 8, 2021 · c5f9f86 · c5f9f86
1 parent 06987b4
commit c5f9f86
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/jax/api.py b/jax/api.py
@@ -2273,6 +2273,8 @@ def __eq__(self, other):
   def __hash__(self):
     return hash((self.shape, self.dtype))
 
+
+@overrideable('eval_shape')
 def eval_shape(fun: Callable, *args, **kwargs):
   """Compute the shape/dtype of ``fun`` without any FLOPs.
 
@@ -2623,6 +2625,7 @@ def vjpfun(ct):
     return ans, vjpfun
   defvjp_all(fun, custom_vjp)
 
+@overrideable('invertible')
 def invertible(fun: Callable) -> Callable:
   """Asserts that the decorated function is invertible.