How to modify a global variable in jit without returning it as a return value?

[sihengz02]

I'm so sorry to ask a question which I think the answer should be easy, but I can't figure it out.
If I have a forward(params, data) function and it does some operations and return the loss, so to jit it, I can't have any side-effect operations in this forward function. But now I need to modify some global variables, like x or something, so I think I can use functional programming and modify the forward function to be forward(params, x, data) and return loss and the modified x just like:

def forward(params, x, data):
    ......
    return loss, x

But if I do this, I can't calculate the gradient of the function forward because it no more return only a single scalar loss.
So how to fix this problem? Thanks for your answering!

[jakevdp]

You can do this using the has_aux argument to jax.grad, which lets you compute the gradient of a function that returns its value along with auxiliary data. For example:

from jax import grad

def f(x, y):
  return y * x, y + 1

grad_f = grad(f, has_aux=True)

print(*grad_f(2.0, 3))
# 3.0, 4