| 23 | |
| 24 | |
| 25 | class Arange(ArrayExpr): |
| 26 | _parameters = ["start", "stop", "step", "chunks", "like", "dtype", "kwargs"] |
| 27 | _defaults = {"chunks": "auto", "like": None, "dtype": None, "kwargs": None} |
| 28 | |
| 29 | @functools.cached_property |
| 30 | def num_rows(self): |
| 31 | return int(max(np.ceil((self.stop - self.start) / self.step), 0)) |
| 32 | |
| 33 | @functools.cached_property |
| 34 | def dtype(self): |
| 35 | return ( |
| 36 | self.operand("dtype") |
| 37 | or np.arange( |
| 38 | self.start, |
| 39 | self.stop, |
| 40 | self.step * self.num_rows if self.num_rows else self.step, |
| 41 | ).dtype |
| 42 | ) |
| 43 | |
| 44 | @functools.cached_property |
| 45 | def _meta(self): |
| 46 | return meta_from_array(self.like, ndim=1, dtype=self.dtype) |
| 47 | |
| 48 | @functools.cached_property |
| 49 | def chunks(self): |
| 50 | return normalize_chunks( |
| 51 | self.operand("chunks"), (self.num_rows,), dtype=self.dtype |
| 52 | ) |
| 53 | |
| 54 | def _layer(self) -> dict: |
| 55 | dsk = {} |
| 56 | elem_count = 0 |
| 57 | start, step = self.start, self.step |
| 58 | like = self.like |
| 59 | func = partial(_arange, like=like) |
| 60 | |
| 61 | for i, bs in enumerate(self.chunks[0]): |
| 62 | blockstart = start + (elem_count * step) |
| 63 | blockstop = start + ((elem_count + bs) * step) |
| 64 | task = Task( |
| 65 | (self._name, i), |
| 66 | func, |
| 67 | blockstart, |
| 68 | blockstop, |
| 69 | step, |
| 70 | bs, |
| 71 | self.dtype, |
| 72 | ) |
| 73 | dsk[(self._name, i)] = task |
| 74 | elem_count += bs |
| 75 | return dsk |
| 76 | |
| 77 | |
| 78 | class Linspace(Arange): |
no outgoing calls
no test coverage detected
searching dependent graphs…