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