def test_broadcast_like_3dim(test_case):
input = flow.Tensor(np.ones(shape=(1, 3, 2), dtype=np.float32),
dtype=flow.float32)
like_tensor = flow.Tensor(np.ones(shape=(3, 3, 2), dtype=np.float32),
dtype=flow.float32)
of_out = flow.broadcast_like(input, like_tensor, broadcast_axes=(0, ))
np_out = np.ones(shape=(3, 3, 2))
test_case.assertTrue(np.allclose(of_out.numpy(), np_out))
def _test_broadcast_like_4dim(test_case, device):
input = flow.Tensor(
np.ones(shape=(1, 3, 2, 1), dtype=np.float32),
dtype=flow.float32,
device=flow.device(device),
)
like_tensor = flow.Tensor(
np.ones(shape=(3, 3, 2, 3), dtype=np.float32),
dtype=flow.float32,
device=flow.device(device),
)
of_out = flow.broadcast_like(input, like_tensor, broadcast_axes=(0, 3))
np_out = np.ones(shape=(3, 3, 2, 3))
test_case.assertTrue(np.allclose(of_out.numpy(), np_out, 1e-5, 1e-5))
def _test_broadcast_like_backward(test_case, device):
input = flow.Tensor(
np.ones(shape=(3, 1, 1), dtype=np.float32),
dtype=flow.float32,
device=flow.device(device),
requires_grad=True,
)
like_tensor = flow.Tensor(
np.ones(shape=(3, 3, 3), dtype=np.float32),
dtype=flow.float32,
device=flow.device(device),
requires_grad=True,
)
of_out = flow.broadcast_like(input, like_tensor, broadcast_axes=(1, 2))
of_out = of_out.sum()
of_out.backward()
np_grad = [[[9.0]], [[9.0]], [[9.0]]]
test_case.assertTrue(np.allclose(input.grad.numpy(), np_grad, 1e-5, 1e-5))