Directly use the backbone functional graph for CausalLM.generate()

[mattdangerw]

Currently, to support the extra inputs we need for generation (e.g. cache, index, encoder hidden states for seq2seq), we are using layers from our backbone class while disregarding the functional graph and layer connectivity of the backbone. See call_with_cache. If we were able to directly use the backbone graph for generation, we would support a lot more advanced generative use cases.

Keras recently added support for optional functional inputs. We should build on that by adding a number of optional inputs to our backbones (e.g. cache, cache_index, token_positions, attention_mask). This would allow customization in a lot of directions:

• Backbones would be more readily useful for more advance non-generative cases without needing to reach into sublayers.
• Generation would be more easily customizable by passing a modified backbone to a CausalLM.