sirhcm/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

remove UOp.contiguous override [PR] (#17012)

Browse Source 
also cleaned up max_shard_shape
This commit is contained in:
chenyu
2026-07-13 16:10:58 -04:00
committed by GitHub
parent 681a5e0cfd
commit c898dfe150
2 changed files with 4 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/mixin/elementwise.py
+3 -1
View File
@@ -49,7 +49,9 @@ class ElementwiseMixin(CreationMixin):
"""
Returns a contiguous tensor.
"""
return self._wrap_uop(self._uop.contiguous(**kwargs))
uop = self._uop
if uop.op is Ops.CONTIGUOUS or self.device is None or uop.has_buffer_identity(): return self._wrap_uop(uop)
return self._wrap_uop(uop.alu(Ops.CONTIGUOUS, **kwargs))
def contiguous_backward(self) -> Self:
"""
tinygrad/uop/ops.py
+1 -8
View File
@@ -446,9 +446,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
return tuple(x//len(self.device) if i == self.axis else x for i,x in enumerate(self.shape))
@property
def max_shard_shape(self) -> tuple[int, ...]:
if not isinstance(self.device, tuple) or self.axis is None: return self.max_shape
return tuple(x//len(self.device) if i == self.axis else x for i,x in enumerate(self.max_shape))
def max_shard_shape(self) -> tuple[int, ...]: return to_max_shape(self.shard_shape)
@functools.cached_property
def ended_ranges(self) -> tuple[UOp, ...]:
@@ -622,11 +620,6 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
if isinstance(arg, Ops): arg = (arg, 0)
return UOp(Ops.REDUCE, src=(self,)+src, arg=arg, **kwargs)
def contiguous(self, *args, **kwargs):
if self.op is Ops.CONTIGUOUS: return self
if self.device is None: return self
if self.has_buffer_identity(): return self
return UOp(Ops.CONTIGUOUS, src=(self,)+args, **kwargs)
def bufferize(self, *args, **kwargs): return UOp(Ops.STAGE, src=(self,)+args, **kwargs)
def allreduce(self, op, device:str|tuple[str, ...]):
assert isinstance(self.device, tuple), f"allreduce must be on tuple {self.device} isn't"