pystencils-autodiff 0.3.0

Automatic differentiation and generation of Torch/Tensorflow operations with pystencils (https://i10git.cs.fau.de/pycodegen/pystencils)

pystencils_autodiff

This repo adds automatic differentiation to pystencils.

Installation

Install via pip:

pip install pystencils-autodiff

or if you downloaded this repository using:

pip install -e .

Then, you can access the submodule pystencils.autodiff.

import pystencils.autodiff

Usage

Create a pystencils.AssignmentCollection with pystencils:

import sympy
import pystencils

z, y, x = pystencils.fields("z, y, x: [20,30]")

forward_assignments = pystencils.AssignmentCollection({
    z[0, 0]: x[0, 0] * sympy.log(x[0, 0] * y[0, 0])
})

print(forward_assignments)

Subexpressions:
Main Assignments:
     z[0,0] ← x_C*log(x_C*y_C)

You can then obtain the corresponding backward assignments:

from pystencils.autodiff import AutoDiffOp, create_backward_assignments
backward_assignments = create_backward_assignments(forward_assignments)

print(backward_assignments)

You can see the derivatives with respective to the two inputs multiplied by the gradient diffz_C of the output z_C.

Subexpressions:
Main Assignments:
    \hat{x}[0,0] ← diffz_C*(log(x_C*y_C) + 1)
    \hat{y}[0,0] ← diffz_C*x_C/y_C

You can also use the class AutoDiffOp to obtain both the assignments (if you are curious) and auto-differentiable operations for Tensorflow…

op = AutoDiffOp(forward_assignments)
backward_assignments = op.backward_assignments

x_tensor = pystencils.autodiff.tf_variable_from_field(x)
y_tensor = pystencils.autodiff.tf_variable_from_field(y)
tensorflow_op = op.create_tensorflow_op({x: x_tensor, y: y_tensor}, backend='tensorflow')

… or Torch:

x_tensor = pystencils.autodiff.torch_tensor_from_field(x, cuda=False, requires_grad=True)
y_tensor = pystencils.autodiff.torch_tensor_from_field(y, cuda=False, requires_grad=True)

z_tensor = op.create_tensorflow_op({x: x_tensor, y: y_tensor}, backend='torch')

Test Report and Coverage

• https://i10git.cs.fau.de/pycodegen/pystencils_autodiff/-/jobs/299120/artifacts/browse