diff --git a/src/pystencils_autodiff/__init__.py b/src/pystencils_autodiff/__init__.py index 752dd6af62aff1168f9667e2e655e8d5decbb0d9..7c6f4b9fa2650d3c658af2babb1e1099e7d9632c 100644 --- a/src/pystencils_autodiff/__init__.py +++ b/src/pystencils_autodiff/__init__.py @@ -1,3 +1,5 @@ +import sys + import pystencils_autodiff.backends # NOQA from pystencils_autodiff._field_to_tensors import ( # NOQA tf_constant_from_field, tf_placeholder_from_field, tf_scalar_variable_from_field, @@ -15,3 +17,6 @@ __all__ = ['backends', "tf_constant_from_field", " tf_placeholder_from_field", "tf_scalar_variable_from_field", " tf_variable_from_field", "torch_tensor_from_field"] + +sys.modules['pystencils.autodiff'] = pystencils_autodiff +sys.modules['pystencils.autodiff.backends'] = pystencils_autodiff.backends diff --git a/src/pystencils_autodiff/backends/__init__.py b/src/pystencils_autodiff/backends/__init__.py index d2ca875cce957d75f39469e095ea26eabd1c9c94..d22da0fb454f507d5f00117e2f885dc48e2c4884 100644 --- a/src/pystencils_autodiff/backends/__init__.py +++ b/src/pystencils_autodiff/backends/__init__.py @@ -1,2 +1,9 @@ +""" +Backends for operators to support automatic Differentation + +Currently, we can use pystencils' JIT compilation to register +a Torch or a Tensorflow operation or we can compile a static +library to be directly loaded into Torch/Tensorflow. +""" AVAILABLE_BACKENDS = ['tensorflow', 'torch', 'tensorflow_cpp', 'torch_native']