Ошибка компиляции XLA с пользовательским слоем Keras на TPU

Question

Я написал свой собственный слой, и у меня возникла проблема с его обучением. Когда начинается обучение, я получаю следующую ошибку:

Ошибка компиляции: XLA не может определить постоянную форму времени компиляции для выделенного фрагмента: [?, 512,512,3], форма вывода должна быть компиляцией постоянная времени

Мой пользовательский слой выглядит примерно так (предоставит полный слой в пост-скрипте):

class MyCustomLayer(Layer):

    def __init__(self, name=None):
        super(MyCustomLayer, self).__init__(name=name)

    @tf.function
    def call(self, inputs):
        """
        inputs has two tensors : [images, image_dependent_params]
        """
        images, image_dependent_params = inputs

        outputs = tf.zeros_like(images)
        for i in tf.range(tf.shape(images)[0]):
            # compute outputs[i] based on images[i] and image_dependent_params[i]
            # set outputs[i] by doing zero padded concatenation

        return outputs

l oop - это обходной путь для crop_and_resize op, который не поддерживается для TPU.

PS: Если вам интересно, вот полный код для него:

class AttentionLocalization(Layer):

    def __init__(self, name=None):
        super(AttentionLocalization, self).__init__(name=name)

    @tf.function
    def call(self, inputs):
        """
        inputs has two tensors : [images, rois]
        """
        images, rois = inputs
        tx, ty, tl = rois[:,0], rois[:,1], rois[:,2]
        tl = tf.clip_by_value(tl, 1/3, 1)
        tx = tf.clip_by_value(tx, 0, 1 - tl)
        ty = tf.clip_by_value(ty, 0, 1 - tl)

        size = tf.cast(tf.shape(images)[1], tf.float32)

        tx, ty, tl = rois[:,0], rois[:,1], rois[:,2]
        tl = tf.clip_by_value(tl, 1/3, 1)
        tx = tf.clip_by_value(tx, 0, 1 - tl)
        ty = tf.clip_by_value(ty, 0, 1 - tl)

        tl = tf.cast(tl * size, tf.int32)
        tx = tf.cast(tx * size, tf.int32)
        ty = tf.cast(ty * size, tf.int32)

        outputs = tf.zeros_like(images)
        batch_size = tf.shape(images)[0]
        for i in tf.range(batch_size):
            resized = tf.image.resize(images[i][tx[i]:tx[i]+tl[i], ty[i]:ty[i]+tl[i], :], tf.shape(images)[1:3])
            padding = tf.zeros([batch_size - 1 - i, tf.shape(images)[1], tf.shape(images)[2], tf.shape(images)[3]])
            outputs = tf.concat([outputs[:i], [resized], padding], axis=0)

        return outputs

РЕДАКТИРОВАТЬ

Я понял, что мой пользовательский слой работает на CPU, но не на TPU, где я пытаюсь обучить свою модель.