def test_where_scalar(test_case):
x = 0.5
y = 2.0
condition = flow.Tensor(np.array([1]), dtype=flow.int32)
of_out = flow.where(condition, x, y)
np_out = np.array([0.5])
test_case.assertTrue(np.allclose(of_out.numpy(), np_out))
def _test_where_backward(test_case, device):
x = flow.Tensor(
np.array([[-0.4620, 0.3139], [0.3898, -0.7197], [0.0478, -0.1657]]),
dtype=flow.float32,
device=flow.device(device),
requires_grad=True,
)
y = flow.Tensor(
np.ones(shape=(3, 2)),
dtype=flow.float32,
device=flow.device(device),
requires_grad=True,
)
condition = flow.Tensor(np.array([[0, 1], [1, 0], [1, 0]]),
dtype=flow.int32,
device=flow.device(device))
of_out = flow.where(condition, x, y)
of_out = of_out.sum()
of_out.backward()
test_case.assertTrue(
np.allclose(x.grad.numpy(),
condition.numpy() == 1, 1e-5, 1e-5))
test_case.assertTrue(
np.allclose(y.grad.numpy(),
condition.numpy() == 0, 1e-5, 1e-5))
def test_where(test_case):
x = flow.Tensor(
np.array([[-0.4620, 0.3139], [0.3898, -0.7197], [0.0478, -0.1657]]),
dtype=flow.float32,
)
y = flow.Tensor(np.ones(shape=(3, 2)), dtype=flow.float32)
condition = flow.Tensor(np.array([[0, 1], [1, 0], [1, 0]]), dtype=flow.int32)
of_out = flow.where(condition, x, y)
np_out = np.array([[1.0000, 0.3139], [0.3898, 1.0000], [0.0478, 1.0000]])
test_case.assertTrue(np.allclose(of_out.numpy(), np_out))
def _test_where_dim4(test_case, device):
x = flow.Tensor(
np.array([[[[-0.4620, 0.3139], [0.3898, -0.7197], [0.0478,
-0.1657]]]]),
dtype=flow.float32,
device=flow.device(device),
)
y = flow.Tensor(np.ones(shape=(1, 1, 3, 2)),
dtype=flow.float32,
device=flow.device(device))
condition = flow.Tensor(
np.array([[[[0, 1], [1, 0], [1, 0]]]]),
dtype=flow.int32,
device=flow.device(device),
)
of_out = flow.where(condition, x, y)
np_out = np.array([[[[1.0000, 0.3139], [0.3898, 1.0000], [0.0478,
1.0000]]]])
test_case.assertTrue(np.allclose(of_out.numpy(), np_out, 1e-5, 1e-5))