Использование собственного оптимизатора Keras в TensorFlow TPU

Question

Я пытаюсь использовать собственный оптимизатор, созданный мной в Keras, который работает при обычной компиляции и подгонке модели, однако, когда я пытаюсь запустить ее с помощью Google Colab TPU, я сталкиваюсь с ошибкой. Есть ли способ использовать этот оптимизатор без необходимости изучать и переписывать его в чистом синтаксисе TensorFlow?

Для простоты воспроизводимого примера я только что взял код SGD прямо из Keras и создал новый класс с ним:

from keras.optimizers import Optimizer
from keras import backend as K
import numpy as np
if K.backend() == 'tensorflow':
    import tensorflow as tf

class SGD2(Optimizer):
    """Stochastic gradient descent optimizer.
    Includes support for momentum,
    learning rate decay, and Nesterov momentum.
    # Arguments
        learning_rate: float >= 0. Learning rate.
        momentum: float >= 0. Parameter that accelerates SGD
            in the relevant direction and dampens oscillations.
        nesterov: boolean. Whether to apply Nesterov momentum.
    """

    def __init__(self, learning_rate=0.01, momentum=0.,
                 nesterov=False, **kwargs):
        learning_rate = kwargs.pop('lr', learning_rate)
        self.initial_decay = kwargs.pop('decay', 0.0)
        super(SGD2, self).__init__(**kwargs)
        with K.name_scope(self.__class__.__name__):
            self.iterations = K.variable(0, dtype='int64', name='iterations')
            self.learning_rate = K.variable(learning_rate, name='learning_rate')
            self.momentum = K.variable(momentum, name='momentum')
            self.decay = K.variable(self.initial_decay, name='decay')
        self.nesterov = nesterov

    def get_updates(self, loss, params):
        grads = self.get_gradients(loss, params)
        self.updates = [K.update_add(self.iterations, 1)]

        lr = self.learning_rate
        if self.initial_decay > 0:
            lr = lr * (1. / (1. + self.decay * K.cast(self.iterations,
                                                      K.dtype(self.decay))))
        # momentum
        shapes = [K.int_shape(p) for p in params]
        moments = [K.zeros(shape, name='moment_' + str(i))
                   for (i, shape) in enumerate(shapes)]
        self.weights = [self.iterations] + moments
        for p, g, m in zip(params, grads, moments):
            v = self.momentum * m - lr * g  # velocity
            self.updates.append(K.update(m, v))

            if self.nesterov:
                new_p = p + self.momentum * v - lr * g
            else:
                new_p = p + v

            # Apply constraints.
            if getattr(p, 'constraint', None) is not None:
                new_p = p.constraint(new_p)

            self.updates.append(K.update(p, new_p))
        return self.updates

    def get_config(self):
        config = {'learning_rate': float(K.get_value(self.learning_rate)),
                  'momentum': float(K.get_value(self.momentum)),
                  'decay': float(K.get_value(self.decay)),
                  'nesterov': self.nesterov}
        base_config = super(SGD, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

Затем выполняется на TPU с помощью:

import os

resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='grpc://' + os.environ['COLAB_TPU_ADDR'])
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

strategy = tf.distribute.experimental.TPUStrategy(resolver)

with strategy.scope():
    model = create_model()
    model.compile(
        optimizer= SGD2(),
        loss='binary_crossentropy',
        metrics=['accuracy'])

model.fit(X_train,
          Y_train,
          epochs = 20,
          batch_size = 1000,
          validation_data=(X_valid, Y_valid)
          )

Где create_model() просто создает последовательную модель (в моем случае CNN, но это не имеет отношения к вопросу). Это дает ошибку AttributeError: 'tensorflow.python.framework.ops.EagerTensor' object has no attribute 'type', однако, если бы я использовал 'adam' в качестве оптимизатора, он работал бы отлично.

Любая помощь будет очень кстати, так как я не могу найти что-либо об этом в Интернете!