| 593 | # ***** creation helper functions ***** |
| 594 | |
| 595 | def _multi_like(self, fxn, *args, **kwargs) -> Tensor: |
| 596 | dtype = kwargs.pop("dtype", self.dtype) |
| 597 | if kwargs.get("device") is not None: raise RuntimeError("cannot specify `device` on `*_like` of a multi device tensor") |
| 598 | assert isinstance(self.device, tuple), f"_multi_like needs a multi device tensor, got {self.device}" |
| 599 | if self.uop.axis is None: return fxn(self.shape, *args, dtype=dtype, **kwargs).shard(self.device) |
| 600 | stacked = UOp.mstack(*[fxn(self.uop.shard_shape, *args, device=d, dtype=dtype, **kwargs).uop for d in self.device]) |
| 601 | return Tensor(stacked.multi(self.uop.axis)) |
| 602 | |
| 603 | def full_like(self, fill_value:ConstType, dtype=None, device=None) -> Tensor: |
| 604 | """ |