MCPcopy
hub / github.com/tinygrad/tinygrad / _multi_like

Method _multi_like

tinygrad/tensor.py:595–601  ·  view source on GitHub ↗
(self, fxn, *args, **kwargs)

Source from the content-addressed store, hash-verified

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 """

Callers 1

rand_likeMethod · 0.95

Calls 6

TensorClass · 0.85
mstackMethod · 0.80
multiMethod · 0.80
fxnFunction · 0.50
getMethod · 0.45
shardMethod · 0.45

Tested by

no test coverage detected