| 2043 | return Array(graph, out, chunks, meta=meta) |
| 2044 | |
| 2045 | def _vindex(self, key): |
| 2046 | if not isinstance(key, tuple): |
| 2047 | key = (key,) |
| 2048 | if any(k is None for k in key): |
| 2049 | raise IndexError( |
| 2050 | f"vindex does not support indexing with None (np.newaxis), got {key}" |
| 2051 | ) |
| 2052 | if all(isinstance(k, slice) for k in key): |
| 2053 | if all( |
| 2054 | k.indices(d) == slice(0, d).indices(d) for k, d in zip(key, self.shape) |
| 2055 | ): |
| 2056 | return self |
| 2057 | raise IndexError( |
| 2058 | "vindex requires at least one non-slice to vectorize over " |
| 2059 | "when the slices are not over the entire array (i.e, x[:]). " |
| 2060 | f"Use normal slicing instead when only using slices. Got: {key}" |
| 2061 | ) |
| 2062 | elif any(is_dask_collection(k) for k in key): |
| 2063 | if math.prod(self.numblocks) == 1 and len(key) == 1 and self.ndim == 1: |
| 2064 | idxr = key[0] |
| 2065 | # we can broadcast in this case |
| 2066 | return idxr.map_blocks( |
| 2067 | _numpy_vindex, self, dtype=self.dtype, chunks=idxr.chunks |
| 2068 | ) |
| 2069 | else: |
| 2070 | raise IndexError( |
| 2071 | "vindex does not support indexing with dask objects. Call compute " |
| 2072 | f"on the indexer first to get an evalurated array. Got: {key}" |
| 2073 | ) |
| 2074 | return _vindex(self, *key) |
| 2075 | |
| 2076 | @property |
| 2077 | def vindex(self): |