Changes to get transformers Gemma working, as well as bug fixes

[adconner]

Some additions to get Gemma-2b working, as implemented by transformers, as well as some bug fixes in the in-place versions of sum, mean, etc.

One divergence with torch semantics with this change is I make no attempt to support the torch.no_grad context (arising as the implementation of torch._C._set_grad_enabled). In the context of Gemma it was used for pure efficiency reasons to mark tensors independent of the input; this optimization is unneeded in the jax autograd context. I am supposing the heuristic that such efficiency is the typical reason for this context, and torch tensors will usually be explicitly detached when this is required for the correctness of the model (already supported by torch2jax). If we wish to support the no_grad context, it can be done by tracking this state ourselves and wrapping the result of all operators in stop_gradient if on.

[samuela]

hi @adconner, thanks for submitting this! it would be super cool to get gemma working. there's some existing, in-flight PRs to be aware of for some of these issues:

• with torch.no_grad #21 adds support for torch.no_grad
• Automatically add out kwarg & member & inplace_member #19 adds out kwarg support and should help reduce the number of changes needed in this PR

i propose that we get those merged first since there already in the works and then rebase from there

[samuela]

also would be cool to add this to the readme or some kind of examples docs somewhere once it's all working!

[samuela]

ok #19 and #21 are merged! so we should be in a pretty good place after rebasing

[adconner]

I've rebased @samuela

[samuela]

thanks @adconner !

[samuela]

just ooc gemma is already implemented in jax, no? https://github.com/google-deepmind/gemma just out of curiosity what's your use case look like?

[adconner]

My use case involves possibly collaborating with torch users in the future, and potentially wanting to run experiments with other models aside from gemma.

[samuela]

lgtm, thanks @adconner !