Не может соответствовать модели keras, которая использует tf.gather (градиенты не существуют)

Question

Я пытаюсь собрать модель keras в тензорном потоке 2, которая предполагает использование gather для умножения каждого входа на несколько различных параметров модели. Например, в нижнем слое, если входные данные равны [i1, i2], а вес модели равен [w1, w2], тогда выходные значения будут [i1 * w1 + i2 * w2, i1 * w2 + i2 * w1]:

import tensorflow as tf

class GatherLayer(tf.keras.layers.Layer):
    def __init__(self, activation='linear', **kwargs):
        super(GatherLayer, self).__init__(**kwargs)
        self.activation = tf.keras.activations.get(activation)

    def build(self, _):
        self.w = self.add_weight(
                shape=(2,), name='weights', initializer='random_normal', trainable=True)
        self.b = self.add_weight(
                shape=(1,), name='offset', initializer='random_normal', trainable=True)
        self.m = tf.concat(
                [tf.reshape(tf.gather(self.w, ix), (-1, 1)) for ix in [[0, 1], [1, 0]]], axis=-1)

    def call(self, inputs):
        lin = tf.matmul(inputs, self.m) + self.b
        return self.activation(lin)

    def get_config(self):
        config = super(GatherLayer, self).get_config()
        return config

Я могу построить этот слой очень хорошо, и вычислять градиенты выходных данных относительно весов, используя GradientTape, но когда я пытаюсь использовать его в tf.keras.Sequential модели, я получаю предупреждения об отсутствующих градиентах:

WARNING:tensorflow:Gradients do not exist for variables ['sequential/gather_layer/weights:0'] when minimizing the loss.
WARNING:tensorflow:Gradients do not exist for variables ['sequential/gather_layer/weights:0'] when minimizing the loss.

Очень похожая установка, использующая tf.stack терпит неудачу таким же образом. Кажется, что у tenorflow возникают проблемы с вычислением градиентов для этих тензоров в ленивом режиме, но я не могу найти ничего о таком ограничении в документации или где-либо еще в Интернете. У кого-нибудь есть идеи?

Полный пример кода:

def gen_data():
    x = [[0, 0],
         [0, 1],
         [1, 0],
         [1, 1]]
    y = [[1, 0],
         [0, 0],
         [1, 1],
         [0, 1]]
    while True:
        yield (x, y)

def make_dataset():
    return tf.data.Dataset.from_generator(
            gen_data, (tf.float32, tf.float32), output_shapes=([batchsize, 2], [batchsize, 2]))

mdl = tf.keras.Sequential([GatherLayer()])
mdl.compile(optimizer=tf.keras.optimizers.Adam(), loss='categorical_crossentropy')
ds = make_dataset()
mdl.fit(ds, steps_per_epoch=4, epochs=1)