tom:

If you’re in for a hack, you can make good use of closures  to eliminate the administration overhead. If things go well, that will also make it easier to script-enable your functions.


Thanks! Very useful – I like that interface and may start using something like that for new functions in the future. I may take a shot at also trying to add in another level of sugar to handle non-Tensor inputs/outputs with an interface like that.
Also I just played around a bit more with the new Function interface and modified the ReLU example – is it ok to manually instantiate the object so I can set some attributes on it and collect some stats in it that I never need to backprop through? (Arbitrarily using numpy here to make sure it works for non-trivial objects that aren’t tensors)
import torch
import numpy as np

class MyReLU(torch.autograd.Function):
    def __init__(ctx, extra_inputs):
        ctx.extra_inputs = extra_inputs
    
    @staticmethod
    def forward(ctx, input):
        ctx.save_for_backward(input)
        print('extra_inputs:', ctx.extra_inputs)
        ctx.extra = {'some_stats': np.random.randn(10)}
        return input.clamp(min=0)

    @staticmethod
    def backward(ctx, grad_output):
        input, = ctx.saved_tensors
        grad_input = grad_output.clone()
        grad_input[input < 0] = 0
        return grad_input
    
ctx = MyReLU(extra_inputs=[np.random.randn(10)])
x = torch.randn(10, requires_grad=True)
y = MyReLU.forward(ctx, x)
print(ctx.extra)
print(torch.autograd.grad(y.sum(), x))