Есть ли способ создать пользовательский слой Keras, вес которого вы можете посмотреть?

Question

Я пытаюсь создать какой-то «словарь», который представляет собой тензор формы, скажем, (8, 16, 32), который я хочу инкапсулировать в пользовательский слой Keras. Я хочу использовать этот слой, чтобы взять в качестве входных данных два числа из набора данных и использовать их в качестве индексов для поиска в словаре и вывода в форме (32,). По сути, я пытаюсь сделать что-то подобное, но со слоем Keras, чтобы я мог обучить словарь:

import numpy as np
a = np.zeros((8, 16, 32)) # The dictionary 
output = a[5,14,:] #5 and 14 will be the input index numbers from my dataset
output.shape #(32,)

Я не писал никаких пользовательских слоев Keras раньше. Я перешел по по этой ссылке , и вот чем я закончил

class GridLayer(layers.Layer):

    def __init__(self, len, width, depth):
        super(GridLayer, self).__init__()
        self.depth = depth
        w_init = tf.random_normal_initializer()
        self.w = tf.Variable(initial_value=w_init(shape=(len, width, depth),
                                                  dtype='float32'),
                                                  trainable=True)
    def call(self, inputs):
        # The input will be of shape ((None, 2, 1)) 
        return tf.slice(self.w, [inputs[1][0], inputs[1][1], 0], [1,1,self.depth])

Когда я пытаюсь использовать слой

# Create the layer 
grid_len = 8
grid_width = 16
grid_depth = 32
grid_layer = GridLayer(grid_len, grid_width, grid_depth)

# Feed input to the custom layer 
input_layer = Input((2,1))
output_layer = grid_layer(input_layer)
model = Model(inputs = input_layer, outputs = output_layer)

, я закончил с этой ошибкой

Traceback (most recent call last):
  File "lstm_practice.py", line 39, in <module>
    output_layer = grid_layer(input_layer)
  File "C:\Users\s3702199\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow_core\python\keras\engine\base_layer.py", line 773, in __call__
    outputs = call_fn(cast_inputs, *args, **kwargs)
  File "C:\Users\s3702199\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow_core\python\autograph\impl\api.py", line 237, in wrapper
    raise e.ag_error_metadata.to_exception(e)
ValueError: in converted code:

    lstm_practice.py:28 call  *
        return tf.slice(self.w, [inputs[1][0], inputs[1][1], 0], [1,1,self.depth])
    C:\Users\s3702199\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow_core\python\ops\array_ops.py:951 slice
        return gen_array_ops._slice(input_, begin, size, name=name)
    C:\Users\s3702199\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow_core\python\ops\gen_array_ops.py:8451 _slice
        "Slice", input=input, begin=begin, size=size, name=name)
    C:\Users\s3702199\AppData\Local\Programs\Python\Python36\lib\site-packages\tensorflow_core\python\framework\op_def_library.py:486 _apply_op_helper
        (input_name, err))

    ValueError: Tried to convert 'begin' to a tensor and failed. Error: Shapes must be equal rank, but are 1 and 0
        From merging shape 1 with other shapes. for 'grid_layer/Slice/packed' (op: 'Pack') with input shapes: [1], [1], [].

Как использовать пользовательский слой Keras для выполнения такого рода тензорной нарезки? Спасибо!