More TF int dtype fixes
Author: Rocketknight1
Pull request: yes

This PR fixes (hopefully) the last remaining TF int dtype issues.

- [x] Ensure all integer dummy inputs are int32 and add test
- [x] Ensure all serving signatures support all-int32
- [x] Check that this fixes `test_saved_model_creation_extended`
- [x] Add a second serving signature when saving with `save_pretrained` to support both int dtypes

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Update: This PR now updates all serving signatures and dummy inputs to `tf.int64`. All tests like `test_saved_model_creation_extended` now pass.

However, I couldn't find any way to cleanly save both a `tf.int32` and a `tf.int64` signature. The reason is that our serving methods currently get their signatures from their tf.function() decorator. This decorator makes it very difficult to extract the underlying function, and we can't call `get_concrete_function` on it with a signature that doesn't match its existing signature. This meant it was quite hard even to insert a shim in `save()` or `save_pretrained()` to save both signatures!

I think we might just have to encourage people onto `tf.int64` as the standard int dtype when exporting models, and require them to pass their own custom signature/function is they want something else. Cc @gante @ydshieh @amyeroberts . If anyone has any good ideas, I'm open!

@gante that's a really good point - models serialized this way might throw errors if users pass tokens straight from our tokenizers, which doesn't seem very sensible for us to do. Let me keep digging - maybe I can find some hack to export an int32 and int64 serving signature, because that's probably too big of a problem to just leave in the codebase.

(Though right now the situation is that models serve either only int32 or only int64 based on what's in the serving signature, which differs between models, and that might be even worse)

Update (with a bit of a deep dive into the TF internals):


When you save a subclassed `Model` as `SavedModel`, the default signature is the forward pass of the model, with the input shape and dtypes that were used the *first* time the model was called. In the TF internals, what happens is that `model._set_save_spec()` is called when the model is built with those first inputs. This records the 'spec' of those inputs, and that spec is used to create the model trace for the SavedModel.

We used to have a huge number of problems because we build our models by passing tiny dummy inputs to them, which locked in that tiny dummy shape as the `save_spec` of the model. We get around that now by calling `_set_save_spec()` in the `__init__()` of the model, and passing flexible shapes (with `None` dimensions). This mostly works well!

The one downside is that even though we can save a spec with flexible shapes by doing this, there's no way to save a spec with flexible dtypes, and multiple specs aren't supported.

To save multiple traces with a `SavedModel`, you can use the `signatures` argument to `model.save()`. However, note that these can be a bit hard for users to find - if the user calls the model directly, they will only get the trace from the `save_spec`, which is locked to one dtype. In other words, you get behaviour like this (assuming the save_spec was `tf.int64`):

```python
# This works
loaded_model(int64_inputs)

# This complains that it can't find a matching signature
loaded_model(int32_inputs)
```

If you explicitly pass a `tf.int32` signature as well as a `tf.int64` signature to the `signatures` argument of `model.save()`, this doesn't fix the issue above. However, you will be able to do this:

```python
# This works
loaded_model.signatures["serving_int32"](int32_inputs)
```

Ideally, I'd like the `SavedModel` to be directly callable with `int32` or `int64`, but I haven't been able to find any way to make this possible, so unfortunately I think we have to pick a 'preferred' int dtype and support the other one only via an awkward call to `model.signatures`. I've added a commit to support this, but I still wish I could find a better approach.

I believe you re-run the target test `test_saved_model_creation_extended` and maybe a few ones after this change, right?
If so, still good for me 💯 !

Yes, I checked those tests again!