Ошибка MatMul: вход не является матрицей в модели Эдварда

Question

Я пытаюсь реализовать модель иерархической логистической регрессии в Эдварде и сталкиваюсь с проблемами со спецификацией модели. Это полный код

def build_toy_dataset(N, D, groups, noise_std=0.1):
    n_wide = N * D
    X = np.linspace(-6, 6, num=n_wide)
    y = np.tanh(X) + np.random.normal(0, noise_std, size=n_wide)
    y[y < 0.5] = 0
    y[y >= 0.5] = 1
    X = (X - 4.0) / 4.0
    X = X.reshape((N, D))
    groups = np.round(np.random.uniform(low=0, high=groups - 1, size=(N)))
    return X, y, groups

N = 100 # X_train.shape[0]
D = 3 # X_train.shape[1]
T = 5000 # Number of posterior samples
num_groups = 4 # number of groups in the dataset

N = 100 # X_train.shape[0]
D = 3 # X_train.shape[1]
T = 5000 # Number of posterior samples
num_groups = 4 # number of groups in the dataset

# Hyperpriors
loc_w = Normal(loc=tf.zeros(D), scale=3.0 * tf.ones(D))
scale_w = Cauchy(loc=tf.zeros(D), scale=3.0 * tf.ones(D))

loc_b = Normal(loc=tf.zeros([]), scale=3.0 * tf.ones([]))
scale_b = Cauchy(loc=tf.zeros([]), scale=3.0 * tf.ones([]))
# Model variables
x = tf.placeholder(tf.float32, [None, D])
group = tf.placeholder(tf.int32, None)

w = Normal(loc=tf.ones([num_groups, D]) * loc_w, scale=tf.ones([num_groups, D]) * scale_w)
b = Normal(loc=tf.ones([num_groups, 1]) * loc_b, scale=tf.ones([num_groups, 1]) * scale_b)

y = Bernoulli(logits=ed.dot(x, tf.gather(w, group)) + tf.gather(b, group))

# Variational Inference
# Location of weights
qloc_w = Normal(loc=tf.get_variable("qw_loc_loc", D), 
                scale=tf.nn.softplus(tf.get_variable("qw_loc_scale", D))
               )
# Scale of weights
qscale_w = Cauchy(loc=tf.get_variable("qw_scale_loc", D), 
                  scale=tf.nn.softplus(tf.get_variable("qw_scale_scale", D))
                 )
# Location of bias
qloc_b = Normal(loc=tf.get_variable("qb_loc_loc", []), 
                scale=tf.nn.softplus(tf.get_variable("qb_loc_scale", []))
               )
# Scale of bias
qscale_b = Cauchy(loc=tf.get_variable("qb_scale_loc", []), 
                  scale=tf.nn.softplus(tf.get_variable("qb_scale_scale", []))
                 )
# Weights
qw = Normal(loc=tf.get_variable("qw_loc", [num_groups, D]), 
                                scale=tf.nn.softplus(tf.get_variable("qw_scale", [num_groups, D]))
                               )
# Bias
qb = Normal(loc=tf.get_variable("qb_loc", [num_groups, 1]),
            scale=tf.nn.softplus(tf.get_variable("qb_scale", [num_groups, 1]))
           )


inference = ed.KLqp({w: qw, b: qb, 
                    loc_w:qloc_w, scale_w: qscale_w,
                    loc_b: qloc_b, scale_b:qscale_b}, 
                    data={x: X_train, y: y_train, group: g_train})
inference.initialize(n_print=10, n_iter=600)

Ошибка, которую я получаю после запуска логического вывода, говорит о том, что In[1] не является матрицей. По одной только ошибке я не могу определить, где именно она произошла. Предполагается, что это происходит на последнем этапе модели, где вычисляются логиты. Я проверил входы по отдельности, и все они выглядели для меня как матрицы.

[[Node: inference/sample/MatMul = MatMul[T=DT_FLOAT, transpose_a=false, transpose_b=false, _device="/job:localhost/replica:0/task:0/device:GPU:0"](inference/sample/control_dependency, inference/sample/ExpandDims)]]
     [[Node: Neg/_47 = _Recv[client_terminated=false, recv_device="/job:localhost/replica:0/task:0/device:CPU:0", send_device="/job:localhost/replica:0/task:0/device:GPU:0", send_device_incarnation=1, tensor_name="edge_2392_Neg", tensor_type=DT_FLOAT, _device="/job:localhost/replica:0/task:0/device:CPU:0"]()]