Add a mechanism to transform the forward pass on Flax models
### Feature request

There should be some way to apply function transformations to Flax models, while not losing the ability to use things like generation utilities.

### Motivation

JAX's main idea is "composable transformations", but currently there's no good way to apply transformations to Flax models. Currently, to apply `my_cool_transformation` to a model, one needs to do something like:

```python
@my_cool_transformation
def wrapper(params, *args, **kwargs):
    return model(*args, params=params, **kwargs)
```

This works fine for training loops and so on, but there doesn't seem to be a way to do this and still be able to use `.generate()`. The reason this would be beneficial is that one can implement things like quantization and LoRA as function transformations, so it would be cool to not lose generation support when doing so.

### Your contribution

I'd be willing to make a PR, but I think this would probably require some modification to the HuggingFace base classes for Flax models.
cc @gante @sanchit-gandhi <|||||>Hey @davisyoshida - that's a great point, would you want to add a composable transformation on top of the `transformers` Flax model (a standard Python class object), or the Flax nn.Module?

If it's the latter (which I believe is the more typical use case), you can first extract the Flax module from the Flax model:
```python
model = FlaxGPT2ForCausalLM.from_pretrained("gpt2")  # standard python object
module = model.module  # flax nn module
```

And then apply any composable transformations to this module (it behaves in the same way as a pure Flax module).

Note that the signature of the module is not the same as the model - this is something that will be addressed by #22499 / https://github.com/huggingface/transformers/pull/22866

You can read more about the `transformers` Flax design philosophy here: https://github.com/huggingface/transformers/tree/main/examples/research_projects/jax-projects#flax-models-in-transformers<|||||>The issue I'm getting at is that extracting the module like that makes you lose access to all the utilities on the model class. Here's an example:

```python
import jax    
import jax.numpy as jnp    
from transformers import AutoTokenizer, FlaxAutoModelForCausalLM    
    
def wrap_model(model):    
    module = model.module    
    def inner(*args, **kwargs):    
        return module.apply(*args, **kwargs)    
    model.module = inner    
    
model_name = 'gpt2'        
model, params = FlaxAutoModelForCausalLM.from_pretrained(model_name, _do_init=False)        
tokenizer = AutoTokenizer.from_pretrained(model_name)        
    
params = jax.device_put(params, jax.devices('gpu')[0])        
    
inputs = jnp.asarray(tokenizer.encode('Hello there.'))[None]        
outputs = model.generate(inputs, params=params)        
print(tokenizer.decode(outputs.sequences[0]))        
    
# Already not very JAX-y since wrap_model mutates `model`        
wrap_model(model)
# Crashes with error: AttributeError: can't set attribute 'module'

outputs = model.generate(inputs, params=params)        
print(tokenizer.decode(outputs.sequences[0]))        
    
# Ideal pure API:        
# my_wrapped_model = wrap_the_model(model)        
# my_wrapped_model.generate(inputs, params=params)
```

If you just extract the module, is there still some way to use generate?
<|||||>Sorry! `model.module` is a `property`, which just returns `model._module`:
https://github.com/huggingface/transformers/blob/495729427045c7a58e040fa9bf6df81c16f54208/src/transformers/modeling_flax_utils.py#L255-L257

You should be able modify `model._module` to access the module class!

> If you just extract the module, is there still some way to use generate?

The generate method is tied to the `FlaxPreTrainedModel`, i.e. the Python class `model`: https://github.com/huggingface/transformers/blob/495729427045c7a58e040fa9bf6df81c16f54208/src/transformers/modeling_flax_utils.py#L158

So I don't think there's any way to generate with just the module that you extract from the model. What you could try is changing the module itself to **also** inherit from `FlaxGenerationMixin`, such that we can call `module.generate`. Note that we'll also have to implement methods like `prepare_inputs_for_generation`:
https://github.com/huggingface/transformers/blob/495729427045c7a58e040fa9bf6df81c16f54208/src/transformers/models/gpt2/modeling_flax_gpt2.py#L745
And `update_inputs_for_generation` for this to work:
https://github.com/huggingface/transformers/blob/495729427045c7a58e040fa9bf6df81c16f54208/src/transformers/models/gpt2/modeling_flax_gpt2.py#L766
Although I'm not sure whether this is possible since Flax modules are just data classes, so you'll have to experiment and see.

I think the easiest would be to define your `wrap_model` function such that it extracts the `_module`, then applies the composition as required, and finally sets the `model._module` attribute again (despite being not super JAX-y, I think this is the easiest way)<|||||>> and finally sets the model._module attribute again 

Ah right I remember running into this when I tried to make generation from quantized models work. The problem is that it's expected that `module` be a proper Flax module, not just a function. Assigning to `_module` in my code above leads to this:

```python
transformers/models/gpt2/modeling_flax_gpt2.py", line 451, in init_cache
    init_variables = self.module.init(
AttributeError: 'function' object has no attribute 'init'
```

You might think to try wrapping `_module`'s `__call__` method like this:

```python
def wrap_model(model):      
    call_fn = model.module.__call__      
    def inner(*args, **kwargs):      
        return call_fn(*args, **kwargs)      
      
    model._module.__call__ = inner
```

But using this (again in the original code I posted), gives the following:
```python
  File "/home/davis/data/venvs/jax/lib/python3.10/site-packages/transformers/models/gpt2/modeling_flax_gpt2.py", line 749, in prepare_inputs_for_generation                                                                                                                     
    past_key_values = self.init_cache(batch_size, max_length)                                                                                                                                                                                                                   
  File "/home/davis/data/venvs/jax/lib/python3.10/site-packages/transformers/models/gpt2/modeling_flax_gpt2.py", line 451, in init_cache                                                                                                                                        
    init_variables = self.module.init(                                                                                                                                                                                                                                          
line 8, in inner                                                                                                                                                                                          
    return call_fn(*args, **kwargs)                                                                                                     
  File "/home/davis/data/venvs/jax/lib/python3.10/site-packages/transformers/models/gpt2/modeling_flax_gpt2.py", line 703, in __call__  
    outputs = self.transformer(                                     
AttributeError: "FlaxGPT2LMHeadModule" object has no attribute "transformer". If "transformer" is defined in '.setup()', remember these fields are only accessible from inside 'init' or 'apply'. 
```

So none of these methods seem to work. I think this should definitely be possible without needing to re-implement methods like `prepare_inputs_for_generation` which are already implemented on the model class.<|||||>Okay I think I figured out the right thing to do, you have to wrap the module's `apply()` method. I think that requires enough indirection that maybe providing a utility for it would be helpful, and it shouldn't be too hard to do.<|||||>Would you like to contribute such a utility @davisyoshida? Or update the docstrings with a note on how this could be done? Would be a nice addition to make it easier to build on top of `transformers` for JAX/Flax models <|||||>Is overwriting `apply()` on the module actually what you guys would like to do as the recommended solution? I think something a bit cleaner might be adding indirection in between the module and places where the model calls it. That way custom behavior could be inserted without needing to modify the Flax object. I'm not sure exactly what the best way to do that would be though.<|||||>What was the kind of utility you had in mind? Not sure I fully follow from your previous comment how this would look other than wrapping the `apply`? https://github.com/huggingface/transformers/issues/24655#issuecomment-1626189281

Perhaps we could go through one or two proposed solutions and discuss them here before proceeding with a PR? Would be great to discuss a bit how this would look before jumping into new code<|||||>Yeah so the simplest option is just something like:

```python
def wrap_apply(model, wrapper):
    model.module.apply = wrapper(model.module.apply)
```

The downside is that it has side effects (although this is probably hard to avoid with the  non-functional API HF went with), but more importantly you can't get the original behavior back (maybe you just want to apply the transformation for evaluation then get back to training).

Another option would be something like:
```python
# On the model:
def set_apply_wrapper(self, wrapper=None):
  self._apply_wrapper = wrapper
 
@property
def module(self):
    # this proxy should wrap self._module and but call
    # wrapper(self._module.apply) whenever 
    return some_proxy_object 
```
This way if you want to restore the model to its original state you can just set the wrapper to `None`.

I think a more ambitious option which is (IMO) more in line with JAX's philosophy, would be to factor the generation utilities out into pure functions which accept whatever arguments they need (e.g. a callable which maps (params, *args, past_cache) -> logits, and one which initializes the cache), then relegate the mixins to just calling those external functions appropriately.

This would let people use the generation utilities much more flexibly.
<|||||>Thanks for the clear explanation - happy to proceed with a PR for design 2 if that works for you short-term? We can then assess how much additional benefit we'd get from a full JAX generate re-factor, since this would be a rather large undertaking as you've outlined.<|||||>Sounds good, I'm willing to put something together. It might be a month or two since I'm pretty slammed atm.<|||||>Perfect! We can also run it by the Flax authors since they're interested in having Transformers' Flax models work more seamlessly with the JAX/Flax libraries 