symforce.codegen.backends.pytorch.pytorch_code_printer module

class PyTorchCodePrinter(settings=None)[source]

Bases: 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.

Parameters:

settings (T.Optional[T.Mapping[str, T.Any]]) –

known_functions = {'Abs': 'abs', 'Sqrt': 'sqrt', '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', 'ln': 'log', 'log': 'log', 'log10': 'log10', 'log1p': 'log1p', 'log2': 'log2', 'loggamma': 'lgamma', 'sin': 'sin', 'sinh': 'sinh', 'tan': 'tan', 'tanh': 'tanh'}
language = 'Python'
doprint(expr, assign_to=None)[source]

Print the expression as code.

Parameters:

  • expr (Expression) – The expression to be printed.

  • assign_to (Symbol, string, MatrixSymbol, list of strings or Symbols (optional)) – If provided, the printed code will set the expression to a variable or multiple variables with the name or names given in assign_to.

Return type:

str