Тензорный тензорот для неизвестного размера партии

Там может быть очевидное решение, но я еще не нашел его. Я хочу сделать простое умножение, где у меня есть один тензор, который дает мне вид вектора весов, а другой - сложенные тензоры (то же число, что и у весов). Использование tf.tensordot кажется простым, но это не работает для tf.tensordot неизвестного размера.

import collections
import tensorflow as tf
tf.reset_default_graph()

x = tf.placeholder(shape=(None, 4, 1), dtype=tf.float32, name='x')
y_true = tf.placeholder(shape=(None, 4, 1), dtype=tf.float32, name='y_true')

# These are the models that I want to combine
linear_model0 = tf.layers.Dense(units=1, name='linear_model0')
linear_model1 = tf.layers.Dense(units=1, name='linear_model1')
agents = collections.OrderedDict()
agents[0] = linear_model0(x)  # shape (?,4,1)
agents[1] = linear_model1(x)  # shape (?,4,1)
stacked = tf.stack(list(agents.values()), axis=1)  # shape (?,2,4,1)

# This is the model that produces the weights
x_flat = tf.layers.Flatten()(x)
weight_model = tf.layers.Dense(units=2, name='weight_model')
weights = weight_model(x_flat)  # shape: (?,2)

# This is the final output
y_pred = tf.tensordot(weights, stacked, axes = 2, name='y_pred')
# PROBLEM HERE: shape: (4,1) instead of (?,4,1) 

# Running the whole thing
sess = tf.Session()
init = tf.global_variables_initializer()
sess.run(init)

# Example1 (output of shape (1,4,1) expected, got (4,1))
print('model', sess.run(y_pred,
                        {x: [[[1], [2], [3], [4]]]}).shape)

# Example2 (output of (2,4,1) expected, got (4,1))
print('model', sess.run(y_pred,
                        {x: [[[1], [2], [3], [4]], [[1], [2], [3], [4]]]}).shape)

Таким образом, умножение работает, как и ожидалось, для первого ввода, но выполняет только первый, а не пакет входов. Любая помощь?

Подобные вопросы, которые не решили мою проблему:

• препятствия в тензорторе тензорного потока с использованием пакетного умножения