pystencils_autodiff
This repo adds automatic differentiation to pystencils.
Installation
Install via pip:
pip install pystencils-autodiffor if you downloaded this repository using:
pip install -e .Then, you can access the submodule pystencils.autodiff.
import pystencils.autodiffUsage
Create a pystencils.AssignmentCollection with pystencils:
import sympy
import pystencils
z, x, y = 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] ← y_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{y}[0,0] ← diffz_C*(log(x_C*y_C) + 1)
\hat{x}[0,0] ← diffz_C*y_C/x_CYou 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')