[SymForce] PyTorch backend

Just functions, and just scalar/matrix inputs/outputs for now

Reviewers: hayk,bradley,nathan,chao,peter,john-m,cedric
Topic: sym-torch
Relative: sf-preamble
GitOrigin-RevId: 18c8c39b8b1c5f0ccea4ae5de534ab2a3eaeb9cb

Author: aaron-skydio

symforce/codegen/backends/pytorch/README.md

@@ -0,0 +1,7 @@
+***THIS MODULE IS EXPERIMENTAL***
+
+Backend for PyTorch.  This generates Python functions that just call PyTorch ops, i.e. each SymForce op in your expression tree becomes a PyTorch op.
+
+It's possible we could do significantly better than this by generating custom PyTorch ops instead.
+
+This currently only supports vector inputs and outputs, we do not have geo or cam types for PyTorch yet.

symforce/codegen/backends/pytorch/__init__.py

@@ -0,0 +1,157 @@
+# ----------------------------------------------------------------------------
+# SymForce - Copyright 2022, Skydio, Inc.
+# This source code is under the Apache 2.0 license found in the LICENSE file.
+# ----------------------------------------------------------------------------
+from __future__ import annotations
+
+import sympy
+from sympy.printing.codeprinter import CodePrinter
+from sympy.printing.pycode import PythonCodePrinter
+
+from symforce import typing as T
+
+_known_functions_torch = {
+    "Abs": "abs",
+    "acos": "acos",
+    "acosh": "acosh",
+    "asin": "asin",
+    "asinh": "asinh",
+    "atan": "atan",
+    "atan2": "atan2",
+    "atanh": "atanh",
+    "ceiling": "ceil",
+    "cos": "cos",
+    "cosh": "cosh",
+    "erf": "erf",
+    "erfc": "erfc",
+    "exp": "exp",
+    "expm1": "expm1",
+    "floor": "floor",
+    "hypot": "hypot",
+    "loggamma": "lgamma",
+    "log": "log",
+    "ln": "log",
+    "log10": "log10",
+    "log1p": "log1p",
+    "log2": "log2",
+    "sin": "sin",
+    "sinh": "sinh",
+    "Sqrt": "sqrt",
+    "tan": "tan",
+    "tanh": "tanh",
+}
+
+_known_constants_math = {
+    "Exp1": "e",
+    "Pi": "pi",
+    "E": "e",
+    "Infinity": "inf",
+    "NaN": "nan",
+    "ComplexInfinity": "nan",
+}
+
+
+def _print_known_const(self: PyTorchCodePrinter, expr: sympy.Expr) -> str:
+    return f"torch.tensor(math.{_known_constants_math[expr.__class__.__name__]}, **tensor_kwargs)"
+
+
+def _print_known_func(self: PyTorchCodePrinter, expr: sympy.Expr) -> str:
+    name = _known_functions_torch[expr.__class__.__name__]
+    return f"torch.{name}({', '.join(map(self._print, expr.args))})"  # pylint: disable=protected-access
+
+
+class PyTorchCodePrinter(CodePrinter):
+    """
+    Symforce customized code printer for PyTorch. Modifies the Sympy printing
+    behavior for codegen compatibility and efficiency.
+
+    This is more different from PythonCodePrinter than it is similar, so we go mostly from scratch
+    and call some methods from that printer where desired.
+    """
+
+    def _format_code(self, lines: T.List[str]) -> T.List[str]:
+        return lines
+
+    def _print_Mod(self, expr: sympy.Mod) -> str:
+        return f"torch.remainder({self._print(expr.args[0])}, {self._print(expr.args[1])})"
+
+    def _print_sign(self, expr: sympy.sign) -> str:
+        return f"torch.sign({self._print(expr.args[0])})"
+
+    def _print_Pow(
+        self, expr: sympy.Pow, rational: bool = False
+    ) -> str:  # pylint: disable=unused-argument
+        # TODO(aaron): Optimize this?
+        return f"torch.pow({self._print(expr.base)}, {self._print(expr.exp)})"
+
+    def _print_Rational(self, expr: sympy.Rational) -> str:
+        # This is py3-only, need decimal points if we want py2
+        return f"torch.tensor({expr.p}/{expr.q}, **tensor_kwargs)"
+
+    def _print_frac(self, expr: sympy.frac) -> str:
+        return self._print_Mod(sympy.Mod(expr.args[0], 1))
+
+    def _print_Integer(self, expr: sympy.Integer) -> str:
+        """
+        Customizations:
+            * Cast all integers to Tensor
+        """
+        return f"torch.tensor({expr.p}, **tensor_kwargs)"
+
+    def _print_NumberSymbol(self, expr: sympy.Expr) -> str:
+        """
+        Customizations:
+            * Cast all NumberSymbols to Tensor
+        """
+        return f"torch.tensor({super()._print_NumberSymbol(expr)}, **tensor_kwargs)"
+
+    def _print_Zero(self, expr: sympy.Expr) -> str:
+        """
+        Customizations:
+            * Cast Zero to Tensor
+        """
+        return "torch.tensor(0, **tensor_kwargs)"
+
+    def _print_Symbol(self, expr: sympy.Symbol) -> str:
+        name = super()._print_Symbol(expr)
+
+        if name in PythonCodePrinter.reserved_words:
+            raise ValueError(
+                f'This expression includes the symbol "{name}" which is a reserved keyword in Python.'
+            )
+
+        return name
+
+    def _print_Max(self, expr: sympy.Max) -> str:
+        if len(expr.args) == 1:
+            return self._print(expr.args[0])
+        else:
+            from sympy.functions.elementary.miscellaneous import Max
+
+            return f"torch.maximum({self._print(expr.args[0])}, {self._print(Max(*expr.args[1:]))})"
+
+    def _print_Min(self, expr: sympy.Min) -> str:
+        if len(expr.args) == 1:
+            return self._print(expr.args[0])
+        else:
+            from sympy.functions.elementary.miscellaneous import Min
+
+            return f"torch.minimum({self._print(expr.args[0])}, {self._print(Min(*expr.args[1:]))})"
+
+    # NOTE(brad): We type ignore the signature because mypy complains that it
+    # does not match that of the sympy base class CodePrinter. This is because the base class
+    # defines _print_Heaviside with: _print_Heaviside = None (see
+    # https://github.com/sympy/sympy/blob/95f0228c033d27731f8707cdbb5bb672e500847d/sympy/printing/codeprinter.py#L446
+    # ).
+    # Despite this, our signature here matches the signatures of the sympy defined subclasses
+    # of CodePrinter. I don't know of any other way to resolve this issue other than to
+    # to type ignore.
+    def _print_Heaviside(self, expr: "sympy.Heaviside") -> str:  # type: ignore[override]
+        return f"torch.heaviside({self._print(expr)}, values=torch.tensor(1.0, **tensor_kwargs))"
+
+
+for k in _known_functions_torch:
+    setattr(PyTorchCodePrinter, f"_print_{k}", _print_known_func)
+
+for k in _known_constants_math:
+    setattr(PyTorchCodePrinter, f"_print_{k}", _print_known_const)

symforce/codegen/backends/pytorch/pytorch_code_printer.py

@@ -0,0 +1,67 @@
+# ----------------------------------------------------------------------------
+# SymForce - Copyright 2022, Skydio, Inc.
+# This source code is under the Apache 2.0 license found in the LICENSE file.
+# ----------------------------------------------------------------------------
+from dataclasses import dataclass
+from pathlib import Path
+
+from sympy.printing.codeprinter import CodePrinter
+
+from symforce import typing as T
+from symforce.codegen.backends.pytorch import pytorch_code_printer
+from symforce.codegen.codegen_config import CodegenConfig
+
+CURRENT_DIR = Path(__file__).parent
+
+
+@dataclass
+class PyTorchConfig(CodegenConfig):
+    """
+    Code generation config for the PyTorch backend.
+
+    Args:
+        doc_comment_line_prefix: Prefix applied to each line in a docstring
+        line_length: Maximum allowed line length in docstrings; used for formatting docstrings.
+        use_eigen_types: Use eigen_lcm types for vectors instead of lists
+        autoformat: Run a code formatter on the generated code
+        custom_preamble: An optional string to be prepended on the front of the rendered template
+        cse_optimizations: Optimizations argument to pass to sf.cse
+        zero_epsilon_behavior: What should codegen do if a default epsilon is not set?
+    """
+
+    doc_comment_line_prefix: str = ""
+    line_length: int = 100
+    use_eigen_types: bool = False
+
+    @classmethod
+    def backend_name(cls) -> str:
+        return "pytorch"
+
+    @classmethod
+    def template_dir(cls) -> Path:
+        return CURRENT_DIR / "templates"
+
+    def templates_to_render(self, generated_file_name: str) -> T.List[T.Tuple[str, str]]:
+        return [
+            ("function/FUNCTION.py.jinja", f"{generated_file_name}.py"),
+            ("function/__init__.py.jinja", "__init__.py"),
+        ]
+
+    def printer(self) -> CodePrinter:
+        return pytorch_code_printer.PyTorchCodePrinter()
+
+    def format_matrix_accessor(self, key: str, i: int, j: int, *, shape: T.Tuple[int, int]) -> str:
+        PyTorchConfig._assert_indices_in_bounds(i, j, shape)
+        if (shape[0] == 1) ^ (shape[1] == 1):
+            return f"{key}[..., {max(i, j)}]"
+        elif shape[0] == 1 and shape[1] == 1:
+            return key
+        else:
+            return f"{key}[..., {i}, {j}]"
+
+    @staticmethod
+    def format_eigen_lcm_accessor(key: str, i: int) -> str:
+        """
+        Format accessor for eigen_lcm types.
+        """
+        raise NotImplementedError("Can't pass eigen_lcm types to PyTorch functions")

symforce/codegen/backends/pytorch/pytorch_config.py

@@ -0,0 +1,28 @@
+{# ----------------------------------------------------------------------------
+ # SymForce - Copyright 2022, Skydio, Inc.
+ # This source code is under the Apache 2.0 license found in the LICENSE file.
+ # ---------------------------------------------------------------------------- #}
+{%- import "../util/util.jinja" as util with context -%}
+
+# pylint: disable=too-many-locals,too-many-lines,too-many-statements,unused-argument
+
+import math # pylint: disable=unused-import
+import typing as T
+
+import torch
+
+
+class TensorKwargs(T.TypedDict):
+    """
+    TypedDict representing args that will be passed to any torch.tensor calls
+    """
+    device: torch.device
+    dtype: torch.dtype
+
+
+{{ util.function_declaration(spec) }}
+    {% if spec.docstring %}
+    {{ util.print_docstring(spec.docstring) | indent(4) }}
+    {% endif %}
+
+    {{ util.expr_code(spec) }}

symforce/codegen/backends/pytorch/templates/function/FUNCTION.py.jinja

@@ -0,0 +1,5 @@
+{# ----------------------------------------------------------------------------
+ # SymForce - Copyright 2022, Skydio, Inc.
+ # This source code is under the Apache 2.0 license found in the LICENSE file.
+ # ---------------------------------------------------------------------------- #}
+from .{{ spec.name }} import {{ spec.name }}