| 82 | assert result.dtype == dtype |
| 83 | |
| 84 | def test_3d_fill_nonna(self, dtype_fill_out_dtype): |
| 85 | dtype, fill_value, out_dtype = dtype_fill_out_dtype |
| 86 | |
| 87 | data = np.random.default_rng(2).integers(0, 2, (5, 4, 3)).astype(dtype) |
| 88 | indexer = [2, 1, 0, -1] |
| 89 | |
| 90 | result = algos.take_nd(data, indexer, axis=0, fill_value=fill_value) |
| 91 | assert (result[[0, 1, 2], :, :] == data[[2, 1, 0], :, :]).all() |
| 92 | assert (result[3, :, :] == fill_value).all() |
| 93 | assert result.dtype == out_dtype |
| 94 | |
| 95 | result = algos.take_nd(data, indexer, axis=1, fill_value=fill_value) |
| 96 | assert (result[:, [0, 1, 2], :] == data[:, [2, 1, 0], :]).all() |
| 97 | assert (result[:, 3, :] == fill_value).all() |
| 98 | assert result.dtype == out_dtype |
| 99 | |
| 100 | result = algos.take_nd(data, indexer, axis=2, fill_value=fill_value) |
| 101 | assert (result[:, :, [0, 1, 2]] == data[:, :, [2, 1, 0]]).all() |
| 102 | assert (result[:, :, 3] == fill_value).all() |
| 103 | assert result.dtype == out_dtype |
| 104 | |
| 105 | indexer = [2, 1, 0, 1] |
| 106 | result = algos.take_nd(data, indexer, axis=0, fill_value=fill_value) |
| 107 | assert (result[[0, 1, 2, 3], :, :] == data[indexer, :, :]).all() |
| 108 | assert result.dtype == dtype |
| 109 | |
| 110 | result = algos.take_nd(data, indexer, axis=1, fill_value=fill_value) |
| 111 | assert (result[:, [0, 1, 2, 3], :] == data[:, indexer, :]).all() |
| 112 | assert result.dtype == dtype |
| 113 | |
| 114 | result = algos.take_nd(data, indexer, axis=2, fill_value=fill_value) |
| 115 | assert (result[:, :, [0, 1, 2, 3]] == data[:, :, indexer]).all() |
| 116 | assert result.dtype == dtype |
| 117 | |
| 118 | def test_1d_other_dtypes(self): |
| 119 | arr = np.random.default_rng(2).standard_normal(10).astype(np.float32) |