Avoid using tf.tile in embeddings for TF models

[ydshieh]

What does this PR do?

Some TF models use

position_embeds = tf.gather(params=self.position_embeddings, indices=position_ids)
position_embeds = tf.tile(input=position_embeds, multiples=(input_shape[0], 1, 1))

which assume that position_ids has size 1 along batch dimension. If users don't specify position_ids, we create it
(before using it)

if position_ids is None:
    position_ids = tf.expand_dims(
        tf.range(start=past_key_values_length, limit=input_shape[1] + past_key_values_length), axis=0
    )

which will have batch size 1. However, in INPUTS_DOCSTRING, it specifies the shape to be (batch_size, seq_len).
If a user provides a full batch for position_ids (although this is very unlikely), tf.tile shouldn't be used here.

This PR fixes this issue.

Who can review?

@Rocketknight1

[Rocketknight1]

I love it, thank you for doing this! I wonder if there's a reason for using the Add() layers like that originally? It feels very odd.

[Rocketknight1]

Either way, it's a straightforward change and I'm happy to merge as-is, so let me know once you're ready.

[ydshieh]

I love it, thank you for doing this! I wonder if there's a reason for using the Add() layers like that originally? It feels very odd.

I also feel the same, and don't know why Add() is used. I removed it here also because Add() requires the shape to be the same, including batch dim (and won't work after I removed tf.tile).

The PR is ready. I can rebase on master to see if I can make the tests green.

[ydshieh]

Failed tests are irrelevant to this PR. Let me know if you prefer to wait and rebase later.

[Rocketknight1]

No, we're seeing those tests on every PR. I'm happy to merge now - let me know whenever the PR is done!

[ydshieh]

No, we're seeing those tests on every PR. I'm happy to merge now - let me know whenever the PR is done!

It's is done. You can merge. Thanks!

[Rocketknight1]

Done!