with torch.no_grad

[mavenlin]

When using torch.no_grad, an error is raised complaining about torch._C.set_grad_enabled is not implemented for Torchish.

This PR added a global variable to track whether gradient is enabled. Any new Torchish element created under the no_grad context will have its .value wrapped with lax.stop_gradient.

[samuela]

i'm hesitant to add a global variable but this may indeed be the simplest route. at some point we ought to also consider supporting/testing https://docs.pytorch.org/docs/stable/generated/torch.enable_grad.html, https://docs.pytorch.org/docs/stable/generated/torch.autograd.grad_mode.set_grad_enabled.html, and https://docs.pytorch.org/docs/stable/generated/torch.is_grad_enabled.html

[samuela]

let's add a comment here documenting why this exists and what parts of the api it interacts with

[mavenlin]

Replaced with forwarding torch._C._set_grad_enabled to its original implementation, which is better than this global variable.

It is somehow necessary, otherwise torch will complain about no __torch_function__ handler defined for torch._C._set_grad_enabled.

[mavenlin]

i'm hesitant to add a global variable but this may indeed be the simplest route. at some point we ought to also consider supporting/testing https://docs.pytorch.org/docs/stable/generated/torch.enable_grad.html, https://docs.pytorch.org/docs/stable/generated/torch.autograd.grad_mode.set_grad_enabled.html, and https://docs.pytorch.org/docs/stable/generated/torch.is_grad_enabled.html

These are added

[mavenlin]

BTW, I've changed value to be a property and added a setter, in this way, the inplace operations will have correct gradient.

[samuela]

nice! reusing torch's internal thing is still global state but hey at least it's not our global state :)

[samuela]

ci should be fixed after b958825

[mavenlin]

ci should be fixed after b958825

Great, thanks!

[samuela]

it's worth mentioning that

with torch.no_grad():
  jax.grad(t2j(torch.sin))(x)

will be None or zero. it's debatable what should really happen here. i'm fine with merging as is and we can revisit later if necessary