Function _AXWb

tests/python/tensorflow/test_nn.py:23–26

(A, X, W, b)

def _AXWb(A, X, W, b):
    X = tf.matmul(X, W)
    Y = tf.reshape(tf.matmul(A, tf.reshape(X, (X.shape[0], -1))), X.shape)
    return Y + b


@pytest.mark.parametrize("out_dim", [1, 2])