How to write a custom function which performs its own forward and backward.

[bbsun]

How to use djl to write a custom function which implments its forward and backward. That means the gradient compuation is not based on the auto-differentiation and it is compuated manually. In PyToch, we can do this like following, is that possible to do similar thing in dji.

class MyReLU(torch.autograd.Function):
"""
We can implement our own custom autograd Functions by subclassing
torch.autograd.Function and implementing the forward and backward passes
which operate on Tensors.
"""

@staticmethod
def forward(ctx, input):
    """
    In the forward pass we receive a Tensor containing the input and return
    a Tensor containing the output. ctx is a context object that can be used
    to stash information for backward computation. You can cache arbitrary
    objects for use in the backward pass using the ctx.save_for_backward method.
    """
    ctx.save_for_backward(input)
    return input.clamp(min=0)

@staticmethod
def backward(ctx, grad_output):
    """
    In the backward pass we receive a Tensor containing the gradient of the loss
    with respect to the output, and we need to compute the gradient of the loss
    with respect to the input.
    """
    input, = ctx.saved_tensors
    grad_input = grad_output.clone()
    grad_input[input < 0] = 0
    return grad_input

[zachgk]

This feature is not currently available: see #140