tfa.optimizers.MultiOptimizer does not work on TPU

[artitechture]

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): google colab
• TensorFlow version and how it was installed (source or binary): 2.12.0
• TensorFlow-Addons version and how it was installed (source or binary): 0.20.0
• Python version: 3.9
• Is GPU used? (yes/no): TPU used

Describe the bug

tfa.optimizers.MultiOptimizer raise InvalidArgumentError on TPU. The following code runs without a error on GPU.

Code to reproduce the issue

https://colab.research.google.com/gist/artitechture/12ddf883cca36290f9580bbbb804df27/tpu_multioptimizer_cifar10.ipynb

Other info / logs

---------------------------------------------------------------------------
InvalidArgumentError                      Traceback (most recent call last)
<ipython-input-8-521e74c800a4> in <cell line: 1>()
----> 1 model.fit(x_train, y_train, validation_data=(x_test, y_test),
      2           epochs=10, batch_size=16*strategy.num_replicas_in_sync)

1 frames
/usr/local/lib/python3.9/dist-packages/keras/utils/traceback_utils.py in error_handler(*args, **kwargs)
     68             # To get the full stack trace, call:
     69             # `tf.debugging.disable_traceback_filtering()`
---> 70             raise e.with_traceback(filtered_tb) from None
     71         finally:
     72             del filtered_tb

/usr/local/lib/python3.9/dist-packages/tensorflow/python/framework/ops.py in _numpy(self)
   1126       return self._numpy_internal()
   1127     except core._NotOkStatusException as e:  # pylint: disable=protected-access
-> 1128       raise core._status_to_exception(e) from None  # pylint: disable=protected-access
   1129 
   1130   @property

InvalidArgumentError: Missing edge input number: 7