diff --git a/docs/index.rst b/docs/index.rst index d45732a96c448cb224ca1abd2052ca3ac7942914..879d7689acde1eafe5185d57a07496c55c479876 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -1,9 +1,10 @@ =================== -pystencils_autodiff +pystencils-autodiff =================== -This is the documentation of **pystencils_autodiff**. +This is the documentation of **pystencils-autodiff**. +.. include:: ../README.rst Contents diff --git a/src/pystencils_autodiff/autodiff.py b/src/pystencils_autodiff/autodiff.py index a81929b9d906acdb359e34b48ab8eeefe0d3c90f..1195b30e5691937e4994cdbfe4739356f8509bf3 100644 --- a/src/pystencils_autodiff/autodiff.py +++ b/src/pystencils_autodiff/autodiff.py @@ -21,6 +21,11 @@ class DiffModes(str, Enum): class AutoDiffOp: + """ + Class to generate code for forward and backward kernels for automatic + differentiation operators. + """ + _REPR_TEMPLATE = jinja2.Template( """Forward: {{ forward_assignments | indent(4) }} @@ -432,6 +437,9 @@ Backward: def time_constant_fields(self): return self._time_constant_fields + def create_torch_op(self, *args, **kwags): + return self.create_tensorflow_op(*args, backend='torch_native', **kwags) + def create_tensorflow_op(self, inputfield_tensor_dict, forward_loop=None,