def _test_squeeze(test_case, device):
np_arr = np.random.rand(1, 1, 1, 3)
input = flow.Tensor(np_arr, device=flow.device(device))
of_shape = flow.squeeze(input, dim=[1, 2]).numpy().shape
np_shape = (1, 3)
test_case.assertTrue(np.array_equal(of_shape, np_shape))
test_case.assertTrue(
np.allclose(
flow.squeeze(input, dim=[1, 2]).numpy(),
np.squeeze(input.numpy(), axis=(1, 2)),
1e-4,
1e-4,
))
def _test_squeeze_backward(test_case, device):
np_arr = np.random.rand(1, 1, 1, 3)
input = flow.Tensor(np_arr, device=flow.device(device), requires_grad=True,)
y = flow.squeeze(input, dim=1).sum()
y.backward()
np_grad = np.ones((1, 1, 1, 3))
test_case.assertTrue(np.array_equal(input.grad.numpy(), np_grad))
def test_squeeze_int(test_case):
input = flow.Tensor(np.array([[[[1, 1, 1]]]]).astype(np.int32))
of_shape = flow.squeeze(input, 1).numpy().shape
np_shape = (1, 1, 3)
test_case.assertTrue(np.array_equal(of_shape, np_shape))
def _test_squeeze_1d_input(test_case, device):
np_arr = np.random.rand(10)
input = flow.Tensor(np_arr, device=flow.device(device))
output = flow.squeeze(input)
test_case.assertTrue(np.allclose(output.numpy(), np_arr, 1e-5, 1e-5))