Add torch_op helper function

firedrakeproject · Feb 5, 2023 · 422b1d5 · 422b1d5
1 parent f388b88
commit 422b1d5
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 0 deletions.
diff --git a/firedrake/__init__.py b/firedrake/__init__.py
@@ -103,6 +103,7 @@
 from firedrake.ensemble import *
 from firedrake.randomfunctiongen import *
 from firedrake.progress_bar import ProgressBar  # noqa: F401
+from firedrake.external_operators import *
 
 from firedrake.logging import *
 # Set default log level

diff --git a/firedrake/external_operators/neural_networks/ml_backend_coupling.py b/firedrake/external_operators/neural_networks/ml_backend_coupling.py
@@ -29,3 +29,7 @@ def __call__(self, *ω):
         metadata = {'F': self.F, 'V_controls': V_controls, 'V_output': V_output}
         φ = partial(self.custom_operator, metadata)
         return φ(*ω)
+
+
+def torch_op(*args, **kwargs):
+    return HybridOperator(*args, **kwargs)