`forward` is calling a strange attribute in custom autograd.Functon class

Hi all,

Let’s say I have a class called NonLinearProblem which solves a specific type of problem, give different values to its arguments. After instantiating a object from this class, one only needs to pass a numpy array to this object like to do the computation. I have also access to analytically obtained gradients

nlp = NonLinearProblem(constraints, #other_args)
input = np.randn(5, )  # any 1D array

# forward function
nlp.__objective(input)  # returns is a scalar 
# backward function
nlp.__gradient(input)  # returns array with the same shape of input

Now, I want to extend autograd.Function to use this class and its corresponding .__objective and .__gradient in forward and backward respectively.

This is what I have done so far:

# I init `nlp` outside of Function

class TopoptFunction2(autograd.Function):
    @staticmethod
    def forward(ctx, densities: torch.Tensor, nlp):
        
        # tps.setOverallDensities()
        output = nlp.__objective(densities.data.numpy())

        ctx.save_for_backward(densities, nlp)

        return torch.FloatTensor(output)

    @staticmethod
    def backward(ctx, grad_output):
        densities, nlp = ctx.saved_tensors
        # tps.setOverallDensities(densities.data.numpy())
        objective_gradient = nlp.__gradient(densities.data.numpy())
        return torch.FloatTensor(objective_gradient * grad_output), None

I am trying to gradcheck by:

topopt = TopoptFunction2.apply
dummy_in = torch.empty(100, dtype=torch.float64).fill_(0.6)
dummy_in.requires_grad = True

test = autograd.gradcheck(topopt, (dummy_in, nonLinearProblem), eps=1e-6, atol=1e-4)
print(test)

But in forwad, I get a strange error where an attirbute combined with name of class and method is trying to be called which obviously does not exist:

Traceback (most recent call last):
  File "/home/nikan3/diff-voxelfem-optimization-problem/VoxelFEM/python/demo_1.py", line 219, in <module>
    test = autograd.gradcheck(topopt, (dummy_in, nonLinearProblem), eps=1e-6, atol=1e-4)
  File "/home/nikan3/.local/lib/python3.6/site-packages/torch/autograd/gradcheck.py", line 323, in gradcheck
    func_out = func(*tupled_inputs)
  File "/home/nikan3/diff-voxelfem-optimization-problem/VoxelFEM/python/demo_1.py", line 187, in forward
    output = nlp.__objective(densities.data.numpy())
AttributeError: 'cyipopt.problem' object has no attribute '_TopoptFunction2__objective'

As you can see, somehow _[ClassName][attributeName] is being called!

Some further details:

  1. cyipopt.problem is type(nlp)
  2. I have created python bindings for this package and it works fine except in this issue

Thanks for your suggestions and comments
Bests