remove is_unrealized_unmasked_const [pr] (#8644)

This commit is contained in:
qazal
2025-01-16 05:27:47 -05:00
committed by GitHub
parent 00e5979897
commit 81a84aa85a
3 changed files with 3 additions and 13 deletions
-10
View File
@@ -2105,16 +2105,6 @@ class TestConst(unittest.TestCase):
print(a.lazydata)
self.assertTrue(tensor_const_pm.rewrite(a.lazydata))
def test_uop_methods(self):
a = Tensor(1)
self.assertTrue(a.lazydata.is_unrealized_unmasked_const())
a = Tensor.ones((4, 4))
self.assertTrue(a.lazydata.is_unrealized_unmasked_const())
a = Tensor.ones((4, 4)).pad((1, 1),)
self.assertFalse(a.lazydata.is_unrealized_unmasked_const())
def test_const_schedule(self):
a = Tensor.ones((4, 4))
sched = a.schedule()
-1
View File
@@ -447,7 +447,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
# COPY is COPY(DEVICE, copyin.base) -> VIEW(copyin.st)
return UOp(Ops.COPY, self.base.dtype, (UOp(Ops.DEVICE, arg=device), self.base), clone).view(unwrap(self.st))
def clone(self) -> UOp: return self.copy_to_device(self.device, clone=True)
def is_unrealized_unmasked_const(self): return self.base.op is Ops.CONST and all(v.mask is None for v in unwrap(self.st).views)
@property
def lbs(self): return [self]
@property
+3 -2
View File
@@ -3130,9 +3130,10 @@ class Tensor(SimpleMathTrait):
# broadcast
return x._broadcast_to(out_shape:=_broadcast_shape(x.shape, y.shape)), y._broadcast_to(out_shape)
# TODO: tensor should stop checking if things are const
def _to_const_val(self, x:Union[Tensor, ConstType]) -> Union[Tensor, ConstType]:
return x.lazydata.const_arg if isinstance(x, Tensor) and isinstance(x.lazydata, UOp) and x.lazydata.is_unrealized_unmasked_const() \
and not x.requires_grad and self._broadcasted(x)[0].shape == self.shape else x
return x.lazydata.const_arg if isinstance(x, Tensor) and isinstance(x.lazydata, UOp) and x.lazydata.base.op is Ops.CONST \
and unwrap(x.lazydata.st).views[0].mask is None and not x.requires_grad and self._broadcasted(x)[0].shape == self.shape else x
def add(self, x:Union[Tensor, ConstType], reverse=False) -> Tensor:
"""