Repository: keras-team/keras
Issue: Grouped convolutions layers cannot be serialized when using mixed precision

**System information**.
- Have I written custom code (as opposed to using a stock example script provided in Keras): no
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): macOS
- TensorFlow installed from (source or binary): binary
- TensorFlow version (use command below): 2.9.0
- Python version: 3.9

**Describe the problem**.

Models that include grouped convolutions throw many warnings and will fail to serialize the grouped convolutions when saving a model that was constructed using mixed precision:
```
INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 24, 24, 64), dtype=tf.float16, name='inputs'), <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7fa9b9d6ec90>), {}).
```

**Describe the current behavior**.
Currently TensorFlow will log the `Unsupported signature for serialization` warning message 4 times per group convolution layer. This is only a problem when using both grouped convolutions and mixed precision. Using either on its own doesn't lead to any problematic behavior.

**Describe the expected behavior**.
Serializing grouped convolution should work also work correctly when using mixed precision training.

**[Contributing](https://github.com/keras-team/keras/blob/master/CONTRIBUTING.md)**.

- Do you want to contribute a PR? (yes/no): yes
- Briefly describe your candidate solution(if contributing):

This problem was introduced in #15868 (/cc @mattdangerw). The proper fix would be to add support for grouped convolutions to the CPU kernels in TensorFlow (`b/213173659`) which would allow us to remove the problematic workaround.
Alternatively this problem can be fixed by adding support for properly serialising `tf.function` take a `AutoCastVariable` as input or  by explicitly converting the `AutoCastVariable` to a tensor before passing it into the `tf.function`:
```diff
diff --git a/keras/layers/convolutional/base_conv.py b/keras/layers/convolutional/base_conv.py
index ed058e53..54a479be 100644
--- a/keras/layers/convolutional/base_conv.py
+++ b/keras/layers/convolutional/base_conv.py
@@ -273,7 +273,9 @@ class Conv(Layer):
             inputs = tf.pad(inputs, self._compute_causal_padding(inputs))
 
         if self.groups > 1:
-            outputs = self._jit_compiled_convolution_op(inputs, self.kernel)
+            outputs = self._jit_compiled_convolution_op(
+                inputs, tf.convert_to_tensor(self.kernel)
+            )
         else:
             outputs = self.convolution_op(inputs, self.kernel)

 ```
@reedwm What do you think would be the best solution here? Adding `tf.convert_to_tensor` before passing the variable to `_jit_compiled_convolution_op` seems to be the simplest option, or would this have unintended side effects?

**Standalone code to reproduce the issue**.

```python
import tensorflow as tf

tf.keras.mixed_precision.set_global_policy("mixed_float16")

model = tf.keras.Sequential(
    [tf.keras.layers.Conv2D(64, 3, groups=8, padding="same", input_shape=(24, 24, 64))]
)
model.compile("sgd")
model.save("/tmp/test")
```
```
INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 24, 24, 64), dtype=tf.float16, name='inputs'), <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7fa9b9d6ec90>), {}).
INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 24, 24, 64), dtype=tf.float16, name='inputs'), <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7fa9b9bcb890>), {}).
INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 24, 24, 64), dtype=tf.float16, name='inputs'), <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7fa9b9d6ec90>), {}).
INFO:tensorflow:Unsupported signature for serialization: ((TensorSpec(shape=(None, 24, 24, 64), dtype=tf.float16, name='inputs'), <tensorflow.python.framework.func_graph.UnknownArgument object at 0x7fa9b9bcb890>), {}).
WARNING:absl:Found untraced functions such as _jit_compiled_convolution_op while saving (showing 1 of 1). These functions will not be directly callable after loading.
INFO:tensorflow:Assets written to: /tmp/test/assets
INFO:tensorflow:Assets written to: /tmp/test/assets
```

A full reproduction can be found in [this notebook](https://colab.research.google.com/drive/1p4NnkQnmp0VFB5iWGmTvq3J-7as5EuHW?usp=sharing).

Changed files:
keras/layers/convolutional/base_conv.py

diff --git a/keras/layers/convolutional/base_conv.py b/keras/layers/convolutional/base_conv.py
index ed058e53bf..54a479bef2 100644
--- a/keras/layers/convolutional/base_conv.py
+++ b/keras/layers/convolutional/base_conv.py
@@ -273,7 +273,9 @@ class Conv(Layer):
             inputs = tf.pad(inputs, self._compute_causal_padding(inputs))
 
         if self.groups > 1:
-            outputs = self._jit_compiled_convolution_op(inputs, self.kernel)
+            outputs = self._jit_compiled_convolution_op(
+                inputs, tf.convert_to_tensor(self.kernel)
+            )
         else:
             outputs = self.convolution_op(inputs, self.kernel)