(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|set[DType]|None=None,
src:tuple[UPat, ...]|list[UPat]|UPat|None=None, arg:Any=None,
name:str|None=None, allow_any_len:bool=False, custom_early_reject:set[Ops]|None=None, location=None, is_any:bool=False)
| 1157 | class UPat(OpMixin): |
| 1158 | __slots__ = ("op", "match_dtype", "arg", "name", "src", "is_any") |
| 1159 | def __init__(self, op:Ops|tuple[Ops, ...]|set[Ops]|None=None, dtype:DType|tuple[DType, ...]|set[DType]|None=None, |
| 1160 | src:tuple[UPat, ...]|list[UPat]|UPat|None=None, arg:Any=None, |
| 1161 | name:str|None=None, allow_any_len:bool=False, custom_early_reject:set[Ops]|None=None, location=None, is_any:bool=False): |
| 1162 | assert op is None or isinstance(op, (Ops, tuple, set)), f"op must be Ops or tuple of Ops, not {op!r}" |
| 1163 | self.op: tuple[Ops, ...]|None = (op,) if isinstance(op, Ops) else (tuple(op) if isinstance(op, set) else op) |
| 1164 | self.match_dtype: tuple[DType, ...]|None = (dtype,) if isinstance(dtype, DType) else (tuple(dtype) if isinstance(dtype, set) else dtype) |
| 1165 | self.arg, self.name, self._in_src, self.custom_early_reject = arg, name, src, custom_early_reject |
| 1166 | self.src: Any = None |
| 1167 | self.is_any = is_any |
| 1168 | assert self.name != "ctx", "UPat can't be named ctx" |
| 1169 | assert dtype is None or isinstance(dtype, DType) or all(isinstance(x, DType) for x in dtype), f"invalid dtype {dtype}" |
| 1170 | |
| 1171 | # try all permutations if it's a list |
| 1172 | if isinstance(src, list): self.src = list(itertools.permutations(src)) if not all_same(src) else [tuple(src)] |
| 1173 | # only one if it's a tuple |
| 1174 | elif isinstance(src, tuple): self.src = [src] |
| 1175 | # repeat if it's a UPat |
| 1176 | elif isinstance(src, UPat): self.src = [itertools.repeat(src)] |
| 1177 | |
| 1178 | self.strict_length = not (allow_any_len or isinstance(src, UPat) or src is None) |
| 1179 | self.required_len: int = 0 if isinstance(src, UPat) or src is None else len(src) |
| 1180 | self.location = location or get_location() |
| 1181 | |
| 1182 | if custom_early_reject is not None: self.early_reject = custom_early_reject |
| 1183 | else: |
| 1184 | upat_match = [src] if isinstance(src, UPat) else ([] if src is None else self.src[0]) |
| 1185 | self.early_reject = {pp.op[0] for pp in upat_match if pp.op is not None and len(pp.op) == 1} |
| 1186 | |
| 1187 | @property |
| 1188 | def dtype(self) -> DType: return self.match_dtype[0] if self.match_dtype is not None else dtypes.void |
nothing calls this directly
no test coverage detected