def _test_where_x_y_none(test_case, device):
condition = flow.tensor(
np.array([[[-0.462, 0.3139], [0.3898, -0.7197], [0.0478, -0.1657]]]),
dtype=flow.float32,
device=flow.device(device),
requires_grad=True,
)
of_out = flow.where(condition)
of_nonzero = flow.nonzero(condition, as_tuple=True)
for i in range(len(of_out)):
test_case.assertTrue(
np.allclose(of_out[i].numpy(), of_nonzero[i].numpy(), 1e-05,
1e-05))
def _test_nonzero(test_case, shape, as_tuple, device):
np_input = np.random.randn(*shape)
input = flow.tensor(np_input,
dtype=flow.float32,
device=flow.device(device))
of_out = flow.nonzero(input, as_tuple)
np_out = np_nonzero(np_input, as_tuple)
if as_tuple:
test_case.assertTrue(
np.allclose(tuple(x.numpy() for x in of_out), np_out, 0.0001,
0.0001))
else:
test_case.assertTrue(
np.allclose(of_out.numpy(), np_out, 0.0001, 0.0001))
def __init__(self, num_experts, gates):
"""Create a SparseDispatcher."""
self._gates = gates
self._num_experts = num_experts
# sort experts
sorted_experts, index_sorted_experts = flow.nonzero(gates).sort(0)
# drop indices
_, self._expert_index = sorted_experts.split(1, dim=1)
# get according batch index for each expert
self._batch_index = sorted_experts[index_sorted_experts[:, 1], 0]
# calculate num samples that each expert gets
self._part_sizes = list((gates > 0).sum(0).numpy())
# TODO workaround
for i in range(len(self._part_sizes)):
self._part_sizes[i] = self._part_sizes[i].item()
# expand gates to match with self._batch_index
gates_exp = gates[self._batch_index.flatten()]
self._nonzero_gates = flow.gather(gates_exp, 1, self._expert_index)
def _nonzero(self, as_tuple=False):
return flow.nonzero(self, as_tuple)
def where_op(condition, x=None, y=None):
if x is None and y is None:
return flow.nonzero(condition, as_tuple=True)
return flow._C.where(condition, x, y)