invalids is empty [pr] (#16353)

This commit is contained in:
chenyu
2026-05-25 16:11:38 -04:00
committed by GitHub
parent b73d2d17b9
commit cdeb861828
2 changed files with 11 additions and 8 deletions
+5 -2
View File
@@ -17,7 +17,9 @@ pm_transform_unique_const = PatternMatcher([
])
pm_ctx = PatternMatcher([
(UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx),
(UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="x"),
lambda ctx,x: x.replace(src=(UOp(Ops.LUNIQUE, arg=next(ctx[1])), x.src[1])) if x.src[0].arg > ctx[2] else add_to_ctx(ctx,x)),
(UPat(Ops.BIND, name="x"), add_to_ctx),
(UPat((Ops.AFTER, Ops.CONTIGUOUS), name="x"),
lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) and x.op_in_backward_slice_with_self(Ops.BUFFER) else None),
])+pm_transform_unique_const
@@ -46,6 +48,7 @@ class _function(Generic[ReturnType]):
# run it and do surgery later
with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)):
_function.depth += 1
unique_start = next(UOp.unique_num)
ret = self.fxn(*args, **kwargs)
_function.depth -= 1
if isinstance(ret, Tensor):
@@ -65,7 +68,7 @@ class _function(Generic[ReturnType]):
# the BUFFERs that are left are the implicit inputs
num_explicit = len(call_uops)
uret = graph_rewrite(uret, pm_ctx, (call_uops, itertools.count(0)), bottom_up=True, name="get_implicit_inputs")
uret = graph_rewrite(uret, pm_ctx, (call_uops, itertools.count(0), unique_start), bottom_up=True, name="get_implicit_inputs")
name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__
if not self.allow_implicit:
implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER]
+6 -6
View File
@@ -7,7 +7,7 @@ from tinygrad.mixin.reduce import ReduceMixin
from tinygrad.uop import Ops
from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element
from tinygrad.device import canonicalize_device
from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, InvalidType, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
from tinygrad.dtype import ConstType, DType, DTypeLike, InvalidType, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
from tinygrad.helpers import all_int, argfix, ceildiv, flatten, flat_to_grouped, make_tuple, prod, resolve_pool_pads, round_up
if TYPE_CHECKING:
@@ -20,6 +20,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
@staticmethod
def unique_const(fill_value:ConstType, **kwargs): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
@staticmethod
def empty(*shape, **kwargs): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
@staticmethod
def const(dtype, b, device=None): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
@classmethod
@@ -46,13 +48,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
@classmethod
def invalids(cls, *shape, **kwargs) -> Self:
"""
Creates a tensor with the given shape, filled with Invalid.
Creates an anonymous uninitialized buffer with the given shape.
This is an alternative to Tensor.empty when you want an "anonymous" buffer.
Eventually Tensor.empty will be replaced by this.
You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor.
"""
return cls.full(argfix(*shape), Invalid, **kwargs)
return cls.empty(argfix(*shape), **kwargs)
@classmethod
def zeros(cls, *shape, **kwargs) -> Self: