I’m trying to look up embedding vectors for a given set of ids and I would like to average them based on number of occurrences.

I got the following so far without the weighted average.

from tensorflow import keras
test = [[1], [2,3], [0]]
weights = [[1], [0.8, 0.2], [0]]
ragged_zips_test = ragged.constant(test)
ragged_weights = ragged.constant(weights)
out = keras.layers.Embedding(input_dim=1001, output_dim=5)(ragged_zips_test)
print(out)
out2 = keras.layers.GlobalAveragePooling1D()(out)  # how do i do weighted average?
print(out2)

<tf.RaggedTensor [[[-0.0055670254, -0.0012740865, -0.0053808466, 0.029169906, -0.032036223]],
 [[0.02231938, -0.035754584, 0.023083176, 0.03383306, -0.038948428],
  [0.008367479, 0.049461152, -0.041627835, -0.016800452, 0.0041155815]],
 [[-0.0047082305, 0.0023113713, -0.014311492, -0.019569397, 0.023326907]]]>
tf.Tensor(
[[-0.00556703 -0.00127409 -0.00538085  0.02916991 -0.03203622]
 [ 0.01534343  0.00685328 -0.00927233  0.0085163  -0.01741642]
 [-0.00470823  0.00231137 -0.01431149 -0.0195694   0.02332691]], shape=(3, 5), dtype=float32)

GlobalAveragePooling1D() doesn’t take sample weights as an option. Is there another operation that would let me do a weighted average of these ragged tensors? Or should I be converting the ragged tensors to sparse?