XLA compilation error with custom Keras layer on TPU

Question

I've written my own custom layer and I'm having an issue training it. When training begins, I get the following error:

Compilation failure: XLA can't deduce compile time constant output shape for strided slice: [?,512,512,3], output shape must be a compile-time constant

My custom layer looks something like this (will provide full one in post-script):

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

The loop is a workaround for the crop_and_resize op which isn't supported for TPU.

PS: In case you're interested, here is the full code for it:

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

EDIT

I realised that my custom layer does work on CPU, but not on TPU which is where I'm trying to train my model.