def test_outer_indexer_consistency_with_broadcast_indexes_vectorized(): def nonzero(x): if isinstance(x, np.ndarray) and x.dtype.kind == 'b': x = x.nonzero()[0] return x original = np.random.rand(10, 20, 30) v = Variable(['i', 'j', 'k'], original) I = ReturnItem() # test orthogonally applied indexers indexers = [ I[:], 0, -2, I[:3], np.array([0, 1, 2, 3]), np.array([0]), np.arange(10) < 5 ] for i, j, k in itertools.product(indexers, repeat=3): if isinstance(j, np.ndarray) and j.dtype.kind == 'b': # match size j = np.arange(20) < 4 if isinstance(k, np.ndarray) and k.dtype.kind == 'b': k = np.arange(30) < 8 _, expected, new_order = v._broadcast_indexes_vectorized((i, j, k)) expected_data = nputils.NumpyVIndexAdapter(v.data)[expected.tuple] if new_order: old_order = range(len(new_order)) expected_data = np.moveaxis(expected_data, old_order, new_order) outer_index = (nonzero(i), nonzero(j), nonzero(k)) actual = indexing._outer_to_numpy_indexer(outer_index, v.shape) actual_data = v.data[actual] np.testing.assert_array_equal(actual_data, expected_data)
def test_outer_indexer_consistency_with_broadcast_indexes_vectorized(): def nonzero(x): if isinstance(x, np.ndarray) and x.dtype.kind == 'b': x = x.nonzero()[0] return x original = np.random.rand(10, 20, 30) v = Variable(['i', 'j', 'k'], original) I = ReturnItem() # noqa: E741 # allow ambiguous name # test orthogonally applied indexers indexers = [I[:], 0, -2, I[:3], np.array([0, 1, 2, 3]), np.array([0]), np.arange(10) < 5] for i, j, k in itertools.product(indexers, repeat=3): if isinstance(j, np.ndarray) and j.dtype.kind == 'b': # match size j = np.arange(20) < 4 if isinstance(k, np.ndarray) and k.dtype.kind == 'b': k = np.arange(30) < 8 _, expected, new_order = v._broadcast_indexes_vectorized((i, j, k)) expected_data = nputils.NumpyVIndexAdapter(v.data)[expected.tuple] if new_order: old_order = range(len(new_order)) expected_data = np.moveaxis(expected_data, old_order, new_order) outer_index = indexing.OuterIndexer((nonzero(i), nonzero(j), nonzero(k))) actual = indexing._outer_to_numpy_indexer(outer_index, v.shape) actual_data = v.data[actual] np.testing.assert_array_equal(actual_data, expected_data)