Autograd

Autograd#

The torch.autograd module provides support for automatic differentiation in pytorch. In xlnstorch, the autograd system is extended to support LNS operations on LNSTensor objects. This allows for gradients to be computed with respect to LNS operations, enabling the use of LNSTensor in neural networks and other machine learning models.

LNSFunction#

To define a differentiable custom LNS operation, you can subclass xlnstorch.autograd.LNSFunction. This class provides the necessary methods to implement the forward and backward passes for your custom operation and is analogous to PyTorch’s torch.autograd.Function.

class xlnstorch.autograd.LNSFunction(*args, **kwargs)#

Base class for LNS operations that require custom forward and backward methods. This class should be subclassed for specific LNS operations.

Note

When using the LNSFunction class, you should pass the LNSTensor objects as inputs to the apply method, not their internal values. This allows the autograd system to correctly track the operations and compute gradients. However, in the forward method, the inputs will be the internal values of the LNSTensor objects, so that you can perform the necessary LNS operations directly. All differentiable outputs should be returned as float64 tensors. For example:

import xlnstorch as xltorch

class MyLNSFunction(xltorch.autograd.LNSFunction):

    @staticmethod
    def forward(x, y): # x and y are float64 tensors
        x_packed, y_packed = x.to(torch.int64), y.to(torch.int64)
        result = f(x_packed, y_packed).to(torch.float64)
        return result # result is a float64 tensor

    @staticmethod
    def setup_context(ctx, inputs, output):
        pass

    @staticmethod
    def backward(ctx, grad_output):
        return grad_output, grad_output

a = xltorch.lnstensor([1.0, 2.0], f=23)
b = xltorch.lnstensor([3.0, 4.0], f=23)
c = MyLNSFunction.apply(a, b) # we pass the LNSTensor objects to apply

Fanout Functions#

In early development of xlnstorch, the autograd system would break when fanout occured. This is where a single LNSTensor is used in multiple operations, so that the same value has multiple paths through the computation graph. To handle and detect this, the following functions were used. This problem has been resolved but the functions are still available for reference. Detecting fanout is equivalent to checking whether the computation graph is a tree or not.

has_fanout(root)

Determines if the autograd graph starting from root has any fan-out nodes.

find_fanout(root)

Detect every node in the autograd graph starting from root that has multiple incoming edges (i.e., multiple parents and therefore fan-out).

raise_fanout_error(offenders)

Turn the offenders list returned by find_fanout into a human-readable error message and raise a RuntimeError.
Contents