Writing Your Own Operations#
 Let’s write our own “multiply” operation. There are two components to doing this:
Defining an operation class (a subclass of
Operation
)Writing a function that ultimately calls
mygrad.execute_op(YourOp, ...)
import numpy as np
import mygrad as mg
from mygrad import execute_op
from mygrad.operation_base import Operation
from mygrad.typing import ArrayLike
# All operations should inherit from Operation, or one of its subclasses
class CustomMultiply(Operation):
""" Performs f(x, y) = x * y """
def __call__(self, x: mg.Tensor, y: mg.Tensor) > np.ndarray:
# This method defines the "forward pass" of the operation.
# It must bind the variable tensors to the op and compute
# the output of the operation as a numpy array
# All tensors must be bound as a tuple to the `variables`
# instance variable.
self.variables = (x, y)
# The forward pass should be performed using numpy arrays,
# not the tensors themselves.
x_arr = x.data
y_arr = y.data
return x_arr * y_arr
def backward_var(self, grad, index, **kwargs):
"""Given ``grad = dℒ/df``, computes ``∂ℒ/∂x`` and ``∂ℒ/∂y``
``ℒ`` is assumed to be the terminal node from which ``ℒ.backward()`` was
called.
Parameters

grad : numpy.ndarray
The backpropagated total derivative with respect to the present
operation: dℒ/df. This will have the same shape as f, the result
of the forward pass.
index : Literal[0, 1]
The indexlocation of ``var`` in ``self.variables``
Returns

numpy.ndarray
∂ℒ/∂x_{i}
Raises

SkipGradient"""
x, y = self.variables
x_arr = x.data
y_arr = y.data
# The operation need not incorporate specialized logic for
# broadcasting. The appropriate sumreductions will be performed
# by MyGrad's autodiff system.
if index == 0: # backprop through a
return grad * y.data # ∂ℒ/∂x = (∂ℒ/∂f)(∂f/∂x)
elif index == 1: # backprop through b
return grad * x.data # ∂ℒ/∂y = (∂ℒ/∂f)(∂f/∂y)
# Our function stitches together our operation class with the
# operation arguments via `mygrad.prepare_op`
def custom_multiply(x: ArrayLike, y: ArrayLike, constant=None) > mg.Tensor:
# `execute_op` will take care of:
#  casting `x` and `y` to tensors if they are instead arraylikes
#  propagating 'constant' status to the resulting output based on the inputs
#  handling inplace operations (specified via the `out` parameter)
return execute_op(CustomMultiply, x, y, constant=constant)
We can now use our differentiable function!
>>> x = mg.tensor(2.0)
>>> y = mg.tensor([1.0, 2.0, 3.0])
>>> custom_multiply(x, y).backward()
>>> x.grad, y.grad
(array(6.), array([2., 2., 2.]))
Documentation for mygrad.Operation#
Base class for all tensor operations that support backpropagation of gradients. 


Backpropagates the gradient through all of the operation's inputs, which are stored in the tuple self.variables. 

Given 