Flax models should allow `inputs_embeds`

Feature request

Currently, non-Flax models allow inputs_embeds instead of input_ids (e.g., GPT2)

def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        ...
        inputs_embeds: Optional[torch.FloatTensor] = None,
        ...
    ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
    ...

However, Flax models have no such option (input_ids only).

It would be great if Flax models also had this option so that,

Optionally, instead of passing input_ids you can choose to directly pass an embedded representation.

Motivation

This is useful if you want more control over how to convert input_ids indices into associated vectors than the model’s internal embedding lookup matrix. (from the docs)

Additionally, this can be useful for things like tuning “soft-prompts” (e.g., https://aclanthology.org/2021.emnlp-main.243/)

Your contribution

I’m will try to implement this myself, but I haven’t yet found a solution.

[BigRedT]

@sanchit-gandhi here’s the PR you requested. I was actually able to get it to work with minimal modifications to generation_flax_utils.py.

@mattf1n a similar solution might work for GPT-2 as well?

[BigRedT]

I would like to second @mattf1n 's feature request. This would be super useful for vision-language modeling where we often want to feed a concatenation of image and text features into a sequence-to-sequence model. This approach has become quite popular recently. See for example - VL-T5, GPV-1, GPV-2, UnifiedIO. And given that non-Flax models already support this, would be great to have this implemented for Flax models as well for consistency!