def _test_flatten(test_case, device):
m = flow.nn.Flatten()
x = flow.Tensor(32, 2, 5, 5, device=flow.device(device))
flow.nn.init.uniform_(x)
y = m(x)
test_case.assertTrue(y.shape == flow.Size((32, 50)))
test_case.assertTrue(
np.array_equal(y.numpy().flatten(),
x.numpy().flatten()))
y2 = flow.flatten(x, start_dim=2)
test_case.assertTrue(y2.shape == flow.Size((32, 2, 25)))
test_case.assertTrue(
np.array_equal(y2.numpy().flatten(),
x.numpy().flatten()))
y3 = x.flatten(start_dim=1)
test_case.assertTrue(y3.shape == flow.Size((32, 50)))
test_case.assertTrue(
np.array_equal(y3.numpy().flatten(),
x.numpy().flatten()))
y4 = x.flatten(start_dim=1, end_dim=2)
test_case.assertTrue(y4.shape == flow.Size((32, 10, 5)))
test_case.assertTrue(
np.array_equal(y4.numpy().flatten(),
x.numpy().flatten()))
y5 = flow.flatten(x)
test_case.assertTrue(y5.shape == flow.Size((1600, )))
test_case.assertTrue(
np.array_equal(y5.numpy().flatten(),
x.numpy().flatten()))
def test_mirrored_tensor_and_op(test_case):
x1 = flow.Tensor([[1.0, 2.0]])
test_case.assertEqual(x1.dtype, flow.float32)
test_case.assertEqual(x1.shape, flow.Size((1, 2)))
x2 = flow.Tensor([[1.0], [2.0]])
# TODO(Liang Depeng): change to MatMul module
op = (flow.builtin_op("matmul").Input("a").Input("b").Attr(
"transpose_a",
False).Attr("transpose_b",
False).Attr("alpha", float(1.0)).Output("out").Build())
y = op(x1, x2)[0]
test_case.assertTrue(
np.array_equal(y.numpy(), np.array([[5.0]], dtype=np.float32)))