Understanding grad object

[liukidar]

Hello,
I'm currently working on modifying objax to work in my framework and I have a question about how the jax.grad function is called by _DerivativeBase.__call__. By my understanding, the grad function cares only about the first input given and doesn't modify the BaseStates at all. By running a couple of tests I got the following:

By checking the ids of a BaseState inside f_func, the values contained in the function argument and in the object var collection are the same

id(tuple(self.vc.subset(BaseState))[0].value)
139927022189152
id(state_tensors[0])
139927022189152

By removing the line self.vc.subset(BaseState).assign(state_tensors), the behaviour of my program (based on the mnist-tutorial but I added on layer that keep track of how many time it has been called with a BaseState counter) doesn't change.

So my question is, am I missing something? Is passing the BaseStates around necessary (more in general, is it necessary to pass through f_func any argument if I'm not asking to compute the gradient with respect to it)?

Thank you.