beignet.func.space

Genentech · Apr 22, 2024 · 2a8f5ed · 2a8f5ed
1 parent b9c6c0a
commit 2a8f5ed
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/beignet/func/_space.py b/src/beignet/func/_space.py
@@ -297,7 +297,7 @@ def shift_fn(a: Tensor, b: Tensor, **_) -> Tensor:
 
         return displacement_fn, shift_fn
 
-    def _displacement_fn(
+    def displacement_fn(
         a: Tensor,
         b: Tensor,
         *,
@@ -328,4 +328,4 @@ def shift_fn(a: Tensor, b: Tensor, **_) -> Tensor:
         def shift_fn(a: Tensor, b: Tensor, **_) -> Tensor:
             return a + b
 
-    return _displacement_fn, shift_fn
+    return displacement_fn, shift_fn