finch

AssemblyContext

Bases: Context

Compiles Finch Notation to Finch Assembly. Holds the state of the compilation process.

Source code in src/finch/compile/lower.py

class AssemblyContext(Context):
    """
    Compiles Finch Notation to Finch Assembly. Holds the state of the
    compilation process.
    """

    def __init__(
        self,
        namespace=None,
        preamble=None,
        epilogue=None,
        bindings=None,
        slots=None,
        access_modes=None,
        types=None,
        func_state: HaltState | None = None,
        mode: CompilerMode | None = None,
    ):
        super().__init__(
            namespace=namespace, preamble=preamble, epilogue=epilogue, mode=mode
        )
        if bindings is None:
            bindings = ScopedDict()
        if slots is None:
            slots = ScopedDict()
        if access_modes is None:
            access_modes = ScopedDict()
        if types is None:
            types = ScopedDict()
        self.bindings = bindings
        self.slots = slots
        self.access_modes = access_modes
        self.types = types
        self.func_state = func_state

    def _slot_expr(self, slot):
        match slot:
            case ntn.Fiber(lvl, _):
                return self(lvl.root)
            case _:
                raise TypeError(f"Unrecognized slot state: {slot}")

    @staticmethod
    def _slot_cursor(slot, type_):
        fields = dict(type_.struct_fields)
        if "pos" in fields:
            pos_t = fields["pos"]
        elif "val" in fields and hasattr(fields["val"], "length_type"):
            pos_t = fields["val"].length_type
        elif getattr(type_, "shape_type", None) == ():
            # 0-d tensors store their value directly rather than in a
            # positioned buffer; the cursor position is unused.
            pos_t = np.intp
        else:
            raise TypeError(f"Cannot create an assembly cursor for {type_}")
        return ntn.Fiber(ntn.Root(ntn.Value(slot, type_)), ntn.Literal(pos_t(0)))

    def block(self):
        """
        Create a new block. Preambles and epilogues will stay within this block.
        This is used to create a new context for compiling a block of code.
        """
        blk = super().block()
        blk.bindings = self.bindings
        blk.slots = self.slots
        blk.access_modes = self.access_modes
        blk.types = self.types
        blk.func_state = self.func_state
        return blk

    def scope(self):
        """
        Create a new scoped context that inherits from this one.
        """
        blk = self.block()
        blk.bindings = self.bindings.scope()
        blk.slots = self.slots.scope()
        blk.access_modes = self.access_modes.scope()
        blk.types = self.types.scope()
        return blk

    def should_halt(self):
        """
        Check if the current function should halt.
        This is used to determine if the function has returned.
        """
        assert self.func_state is not None
        return self.func_state.has_returned

    def emit(self):
        return self.preamble + self.epilogue

    def resolve(self, node):
        match node:
            case ntn.Slot(var_n, var_t):
                if var_n in self.slots:
                    view = self.slots[var_n]
                    if view.result_type != var_t:
                        raise TypeError(
                            f"Slot '{var_n}' is declared as type {view.result_type}, "
                            f"but used as type {var_t}."
                        )
                    if isinstance(view, ntn.Fiber):
                        return view
                    raise TypeError(f"Unrecognized slot state: {view}")
                raise KeyError(f"Slot {var_n} not found in context")
            case ntn.Fiber():
                return node
            case ntn.Value(_, _) | ntn.Full():
                return node
            case _:
                raise ValueError(
                    f"Expected Slot, Fiber, Value, or Full, got: {type(node)}"
                )

    def _freeze_tensor(self, tns_var: str, op: ntn.Literal | None) -> None:
        if op is None:
            assert tns_var not in self.access_modes
        else:
            assert self.access_modes[tns_var] == ntn.Update(op)
        self.access_modes[tns_var] = ntn.Read()

    def _thaw_tensor(self, tns_var: str, op: ntn.Literal) -> None:
        assert self.access_modes[tns_var] == ntn.Read()
        self.access_modes[tns_var] = ntn.Update(op)

    def _rm_tensor_from_accesses(self, tns_var: str) -> None:
        del self.access_modes[tns_var]

    # TODO: Move to .pyi file

    @overload
    def __call__(self, prgm: ntn.Module) -> asm.Module:
        pass

    @overload
    def __call__(self, prgm: ntn.Literal) -> asm.Literal:
        pass

    @overload
    def __call__(self, prgm: ntn.Variable) -> asm.Variable:
        pass

    @overload
    def __call__(self, prgm: ntn.NotationExpression) -> asm.AssemblyExpression:
        pass

    @overload
    def __call__(self, prgm: ntn.NotationStatement) -> None:
        pass

    def __call__(self, prgm):
        """
        Lower Finch Notation to Finch Assembly. First we check for early
        simplifications, then we call the normal lowering for the outermost
        node.
        """
        match prgm:
            case ntn.Literal(value):
                return asm.Literal(value)
            case ntn.Value(expr, _):
                return expr
            case ntn.Root(tns):
                return tns.result_type.get_child(self(tns), "lvl")
            case ntn.Child(parent, attr):
                return parent.result_type.level_get_child(self(parent), attr)
            case ntn.Call(f, args):
                f_e = self(f)
                args_e = tuple(self(arg) for arg in args)
                return asm.Call(f_e, args_e)
            case ntn.Assign(var, val):
                self.exec(asm.Assign(self(var), self(val)))
                return None
            case ntn.Variable(var_n, var_t):
                return asm.Variable(var_n, var_t)
            case ntn.Slot(var_n, var_t):
                if var_n in self.types:
                    def_t = self.types[var_n]
                    if def_t != var_t:
                        raise TypeError(
                            f"Slot '{var_n}' is declared as type {def_t}, "
                            f"but used as type {var_t}."
                        )
                if var_n in self.slots:
                    return self._slot_expr(self.slots[var_n])
                raise KeyError(f"Slot '{var_n}' is not defined in the current context.")
            case ntn.Fiber(lvl, _):
                return self(lvl.root)
            case ntn.Unpack(ntn.Slot(var_n, var_t), val):
                val_code = self(val)
                if val.result_type != var_t:
                    raise TypeError(f"Type mismatch: {val.result_type} != {var_t}")
                if var_n in self.slots:
                    raise KeyError(
                        f"Slot {var_n} already exists in context, cannot unpack"
                    )
                if var_n in self.types:
                    raise KeyError(
                        f"Variable '{var_n}' is already defined in the current"
                        f" context, cannot overwrite with slot."
                    )
                slot = asm.Slot(var_n, var_t)
                self.exec(asm.Unpack(slot, val_code))
                self.types[var_n] = var_t
                self._freeze_tensor(var_n, op=None)
                self.slots[var_n] = self._slot_cursor(slot, var_t)
                return None
            case ntn.Repack(ntn.Slot(var_n, var_t), _):
                if var_n not in self.slots or var_n not in self.types:
                    raise KeyError(f"Slot {var_n} not found in context, cannot repack")
                if var_t != self.types[var_n]:
                    raise TypeError(f"Type mismatch: {var_t} != {self.types[var_n]}")
                self._rm_tensor_from_accesses(var_n)
                self.exec(asm.Repack(asm.Slot(var_n, var_t)))
                return None
            case ntn.Unwrap(ntn.Full(val, ())):
                return self(val)
            case ntn.Unwrap(ntn.Access(ntn.Full(val, ()), ntn.Read(), ())):
                return self(val)
            case ntn.Dimension(ntn.Full(_, shape), ntn.Literal(r)):
                return self(shape[r])
            case ntn.Unwrap(ntn.Access(tns, mode, idxs)):
                assert isinstance(mode, ntn.Read)
                assert idxs == ()
                tns = self.resolve(tns)
                return tns.result_type.lower_unwrap(self, tns)
            case ntn.Increment(ntn.Access(tns, mode, idxs), val):
                assert isinstance(mode, ntn.Update)
                assert idxs == ()
                tns = self.resolve(tns)
                op = mode.op
                return tns.result_type.lower_increment(self, tns, op, val)
            case ntn.Block(bodies):
                for body in bodies:
                    self(body)
                return None
            case ntn.Loop(idx, ext, body):
                ext.result_type.lower_loop(
                    self, idx, SymbolicExtent.from_notation(ext), body
                )
                return None
            case ntn.Dimension(tns, ntn.Literal(r)):
                assert isinstance(r, int)
                tns = self.resolve(tns)
                return tns.result_type.lower_dim(self, tns, r)
            case ntn.FillValue(tns):
                tns = self.resolve(tns)
                return tns.result_type.lower_fill(self, tns)
            case ntn.Cached(arg, _):
                return self(arg)
            case ntn.Declare(tns, init, op, shape):
                self._thaw_tensor(tns.name, op)
                tns = self.resolve(tns)
                # A declaration may change a dynamic fill, but not a static one.
                fill = tns.result_type.fill_type
                if is_static(fill) and not (
                    isinstance(init, ntn.Literal)
                    and bool(ffuncs.same(init.val, fill.value))
                ):
                    raise ValueError(
                        f"Declaring {tns} with {init} would change its static "
                        f"fill {fill.value!r}"
                    )
                init_e = self(init)
                op_e = self(op)
                shape_e = [self(s) for s in shape]
                return tns.result_type.lower_declare(self, tns, init_e, op_e, shape_e)
            case ntn.Freeze(tns, op):
                self._freeze_tensor(tns.name, op)
                tns = self.resolve(tns)
                op_e = self(op)
                return tns.result_type.lower_freeze(self, tns, op_e)
            case ntn.Thaw(tns, op):
                self._thaw_tensor(tns.name, op)
                tns = self.resolve(tns)
                op_e = self(op)
                return tns.result_type.lower_thaw(self, tns, op_e)
            case ntn.If(cond, body):
                cond_e = self(cond)
                ctx_2 = self.scope()
                ctx_2(body)
                self.exec(asm.If(cond_e, asm.Block(ctx_2.emit())))
                return None
            case ntn.IfElse(cond, body, else_body):
                cond_e = self(cond)
                ctx_2 = self.scope()
                ctx_2(body)
                ctx_3 = self.scope()
                ctx_3(else_body)
                self.exec(
                    asm.IfElse(cond_e, asm.Block(ctx_2.emit()), asm.Block(ctx_3.emit()))
                )
                return None
            case ntn.Function(
                ntn.Variable(
                    func_n, asm.AssemblyKernelFType(result_type=ret_t) as func_type
                ),
                args,
                body,
            ):
                ctx = self.scope()
                ctx.func_state = HaltState(
                    return_var=asm.Variable(ctx.freshen(f"{func_n}_return"), ret_t)
                )
                blk = ctx.scope()
                blk(body)
                self.exec(
                    asm.Function(
                        asm.Variable(func_n, func_type),
                        tuple(ctx(arg) for arg in args),
                        asm.Block((*blk.emit(), asm.Return(ctx.func_state.return_var))),
                    )
                )
                return None
            case ntn.Return(value):
                if self.func_state is None:
                    raise ValueError("Return statement outside of function.")
                self.exec(asm.Assign(self.func_state.return_var, self(value)))
                return None
            case ntn.Module(funcs):
                ctx = self.scope()
                for func in funcs:
                    ctx(func)
                return asm.Module(ctx.emit())
            case other:
                raise Exception(f"{other} not recognized.")

__call__(prgm)

__call__(prgm: ntn.Module) -> asm.Module

__call__(prgm: ntn.Literal) -> asm.Literal

__call__(prgm: ntn.Variable) -> asm.Variable

__call__(
    prgm: ntn.NotationExpression,
) -> asm.AssemblyExpression

__call__(prgm: ntn.NotationStatement) -> None

Lower Finch Notation to Finch Assembly. First we check for early simplifications, then we call the normal lowering for the outermost node.

Source code in src/finch/compile/lower.py

def __call__(self, prgm):
    """
    Lower Finch Notation to Finch Assembly. First we check for early
    simplifications, then we call the normal lowering for the outermost
    node.
    """
    match prgm:
        case ntn.Literal(value):
            return asm.Literal(value)
        case ntn.Value(expr, _):
            return expr
        case ntn.Root(tns):
            return tns.result_type.get_child(self(tns), "lvl")
        case ntn.Child(parent, attr):
            return parent.result_type.level_get_child(self(parent), attr)
        case ntn.Call(f, args):
            f_e = self(f)
            args_e = tuple(self(arg) for arg in args)
            return asm.Call(f_e, args_e)
        case ntn.Assign(var, val):
            self.exec(asm.Assign(self(var), self(val)))
            return None
        case ntn.Variable(var_n, var_t):
            return asm.Variable(var_n, var_t)
        case ntn.Slot(var_n, var_t):
            if var_n in self.types:
                def_t = self.types[var_n]
                if def_t != var_t:
                    raise TypeError(
                        f"Slot '{var_n}' is declared as type {def_t}, "
                        f"but used as type {var_t}."
                    )
            if var_n in self.slots:
                return self._slot_expr(self.slots[var_n])
            raise KeyError(f"Slot '{var_n}' is not defined in the current context.")
        case ntn.Fiber(lvl, _):
            return self(lvl.root)
        case ntn.Unpack(ntn.Slot(var_n, var_t), val):
            val_code = self(val)
            if val.result_type != var_t:
                raise TypeError(f"Type mismatch: {val.result_type} != {var_t}")
            if var_n in self.slots:
                raise KeyError(
                    f"Slot {var_n} already exists in context, cannot unpack"
                )
            if var_n in self.types:
                raise KeyError(
                    f"Variable '{var_n}' is already defined in the current"
                    f" context, cannot overwrite with slot."
                )
            slot = asm.Slot(var_n, var_t)
            self.exec(asm.Unpack(slot, val_code))
            self.types[var_n] = var_t
            self._freeze_tensor(var_n, op=None)
            self.slots[var_n] = self._slot_cursor(slot, var_t)
            return None
        case ntn.Repack(ntn.Slot(var_n, var_t), _):
            if var_n not in self.slots or var_n not in self.types:
                raise KeyError(f"Slot {var_n} not found in context, cannot repack")
            if var_t != self.types[var_n]:
                raise TypeError(f"Type mismatch: {var_t} != {self.types[var_n]}")
            self._rm_tensor_from_accesses(var_n)
            self.exec(asm.Repack(asm.Slot(var_n, var_t)))
            return None
        case ntn.Unwrap(ntn.Full(val, ())):
            return self(val)
        case ntn.Unwrap(ntn.Access(ntn.Full(val, ()), ntn.Read(), ())):
            return self(val)
        case ntn.Dimension(ntn.Full(_, shape), ntn.Literal(r)):
            return self(shape[r])
        case ntn.Unwrap(ntn.Access(tns, mode, idxs)):
            assert isinstance(mode, ntn.Read)
            assert idxs == ()
            tns = self.resolve(tns)
            return tns.result_type.lower_unwrap(self, tns)
        case ntn.Increment(ntn.Access(tns, mode, idxs), val):
            assert isinstance(mode, ntn.Update)
            assert idxs == ()
            tns = self.resolve(tns)
            op = mode.op
            return tns.result_type.lower_increment(self, tns, op, val)
        case ntn.Block(bodies):
            for body in bodies:
                self(body)
            return None
        case ntn.Loop(idx, ext, body):
            ext.result_type.lower_loop(
                self, idx, SymbolicExtent.from_notation(ext), body
            )
            return None
        case ntn.Dimension(tns, ntn.Literal(r)):
            assert isinstance(r, int)
            tns = self.resolve(tns)
            return tns.result_type.lower_dim(self, tns, r)
        case ntn.FillValue(tns):
            tns = self.resolve(tns)
            return tns.result_type.lower_fill(self, tns)
        case ntn.Cached(arg, _):
            return self(arg)
        case ntn.Declare(tns, init, op, shape):
            self._thaw_tensor(tns.name, op)
            tns = self.resolve(tns)
            # A declaration may change a dynamic fill, but not a static one.
            fill = tns.result_type.fill_type
            if is_static(fill) and not (
                isinstance(init, ntn.Literal)
                and bool(ffuncs.same(init.val, fill.value))
            ):
                raise ValueError(
                    f"Declaring {tns} with {init} would change its static "
                    f"fill {fill.value!r}"
                )
            init_e = self(init)
            op_e = self(op)
            shape_e = [self(s) for s in shape]
            return tns.result_type.lower_declare(self, tns, init_e, op_e, shape_e)
        case ntn.Freeze(tns, op):
            self._freeze_tensor(tns.name, op)
            tns = self.resolve(tns)
            op_e = self(op)
            return tns.result_type.lower_freeze(self, tns, op_e)
        case ntn.Thaw(tns, op):
            self._thaw_tensor(tns.name, op)
            tns = self.resolve(tns)
            op_e = self(op)
            return tns.result_type.lower_thaw(self, tns, op_e)
        case ntn.If(cond, body):
            cond_e = self(cond)
            ctx_2 = self.scope()
            ctx_2(body)
            self.exec(asm.If(cond_e, asm.Block(ctx_2.emit())))
            return None
        case ntn.IfElse(cond, body, else_body):
            cond_e = self(cond)
            ctx_2 = self.scope()
            ctx_2(body)
            ctx_3 = self.scope()
            ctx_3(else_body)
            self.exec(
                asm.IfElse(cond_e, asm.Block(ctx_2.emit()), asm.Block(ctx_3.emit()))
            )
            return None
        case ntn.Function(
            ntn.Variable(
                func_n, asm.AssemblyKernelFType(result_type=ret_t) as func_type
            ),
            args,
            body,
        ):
            ctx = self.scope()
            ctx.func_state = HaltState(
                return_var=asm.Variable(ctx.freshen(f"{func_n}_return"), ret_t)
            )
            blk = ctx.scope()
            blk(body)
            self.exec(
                asm.Function(
                    asm.Variable(func_n, func_type),
                    tuple(ctx(arg) for arg in args),
                    asm.Block((*blk.emit(), asm.Return(ctx.func_state.return_var))),
                )
            )
            return None
        case ntn.Return(value):
            if self.func_state is None:
                raise ValueError("Return statement outside of function.")
            self.exec(asm.Assign(self.func_state.return_var, self(value)))
            return None
        case ntn.Module(funcs):
            ctx = self.scope()
            for func in funcs:
                ctx(func)
            return asm.Module(ctx.emit())
        case other:
            raise Exception(f"{other} not recognized.")

block()

Create a new block. Preambles and epilogues will stay within this block. This is used to create a new context for compiling a block of code.

Source code in src/finch/compile/lower.py

def block(self):
    """
    Create a new block. Preambles and epilogues will stay within this block.
    This is used to create a new context for compiling a block of code.
    """
    blk = super().block()
    blk.bindings = self.bindings
    blk.slots = self.slots
    blk.access_modes = self.access_modes
    blk.types = self.types
    blk.func_state = self.func_state
    return blk

scope()

Create a new scoped context that inherits from this one.

Source code in src/finch/compile/lower.py

def scope(self):
    """
    Create a new scoped context that inherits from this one.
    """
    blk = self.block()
    blk.bindings = self.bindings.scope()
    blk.slots = self.slots.scope()
    blk.access_modes = self.access_modes.scope()
    blk.types = self.types.scope()
    return blk

should_halt()

Check if the current function should halt. This is used to determine if the function has returned.

Source code in src/finch/compile/lower.py

def should_halt(self):
    """
    Check if the current function should halt.
    This is used to determine if the function has returned.
    """
    assert self.func_state is not None
    return self.func_state.has_returned

BufferizedNDArray

Bases: OverrideTensor

Source code in src/finch/tensor/bufferized_ndarray.py

class BufferizedNDArray(OverrideTensor):
    def __init__(
        self,
        val: NumpyBuffer,
        shape: tuple[np.intp, ...],
        strides: tuple[np.intp, ...],
        fill_value: Any = 0,
        device=None,
    ):
        self.val = val
        self._shape = shape
        pos_t = val.ftype.length_type
        self.strides = tuple(pos_t(stride) for stride in strides)
        elem_t = val.ftype.element_type
        self._fill_value = (
            fill_value
            if isinstance(fill_value, AbstractFill)
            else StaticFill(elem_t(fill_value))
        )

        self._device = normalize_device(device)

    def with_fill(self, fill_value: AbstractFill) -> "BufferizedNDArray":
        return BufferizedNDArray(
            self.val,
            self._shape,
            self.strides,
            fill_value,
            device=self._device,
        )

    def to_numpy(self):
        """
        Convert the bufferized NDArray to a NumPy array.
        This is used to get the underlying NumPy array from the bufferized NDArray.
        """
        return self.val.arr.reshape(self._shape, copy=False)

    def to_scipy(self):
        raise NotImplementedError(f"{type(self).__name__} does not support to_scipy.")

    @classmethod
    def from_numpy(
        cls, arr: np.ndarray, fill_value: Any = 0, device=None
    ) -> "BufferizedNDArray":
        if not arr.flags["C_CONTIGUOUS"]:
            arr = np.ascontiguousarray(arr)
        itemsize = arr.dtype.itemsize
        strides = tuple(np.intp(stride // itemsize) for stride in arr.strides)
        shape = tuple(np.intp(s) for s in arr.shape)
        val = NumpyBuffer(arr.reshape(-1, copy=False))
        fill = as_fill(fill_value)
        coerced = np.asarray(fill.value, dtype=arr.dtype).flat[0]
        return BufferizedNDArray(
            val,
            shape,
            strides,
            DynamicFill(coerced) if is_dynamic(fill) else StaticFill(coerced),
            device=device,
        )

    def __array__(self, dtype=None, copy=None):
        out = self.to_numpy()
        if dtype is not None and out.dtype != dtype:
            if copy is not None and not copy:
                raise ValueError(
                    "Unable to avoid copy while creating an array as requested."
                )
            out = out.astype(dtype)
        return out

    @property
    def ftype(self):
        """
        Returns the ftype of the buffer, which is a BufferizedNDArrayFType.
        """
        return BufferizedNDArrayFType(
            buffer_type=ftype(self.val),
            ndim=self.ndim,
            dimension_type=ftype(self._shape),
            fill_type=self._fill_value.ftype,
            device=self._device,
        )

    @property
    def shape(self) -> tuple[int, ...]:
        return tuple(int(s) for s in self._shape)

    @property
    def ndim(self) -> int:
        return len(self._shape)

    @property
    def fill_value(self) -> Any:
        """Default value to fill the tensor."""
        return self._fill_value.value

    @property
    def fill(self) -> AbstractFill:
        return self._fill_value

    @property
    def device(self):
        return self._device

    def to_device(self, device, /, *, stream=None):
        if stream is not None:
            raise ValueError(f"stream argument is not supported; got {stream!r}")
        device = normalize_device(device)
        if device == self.device:
            return self
        return BufferizedNDArray(
            self.val,
            self._shape,
            self.strides,
            self._fill_value,
            device=device,
        )

    @property
    def element_type(self) -> FType:
        """Data type of the tensor elements."""
        return self.ftype.element_type

    @property
    def shape_type(self) -> tuple:
        """Shape type of the tensor."""
        return self.ftype.shape_type

    def declare(self, init, op, shape):
        """
        Declare a bufferized NDArray with the given initialization value,
        operation, and shape.
        """
        if is_dynamic(self._fill_value):
            # Declaring a tensor sets its dynamic fill to the init.
            init = self.val.ftype.element_type(init)
            self._fill_value = DynamicFill(init, self.val.ftype.element_type)
        elif not bool(ffuncs.same(init, self._fill_value.value)):
            raise ValueError(
                f"Declaring with {init!r} would change the static fill "
                f"{self._fill_value.value!r}"
            )
        for i in range(self.val.length()):
            self.val.store(i, init)
        return self

    def freeze(self, op):
        return self

    def thaw(self, op):
        return self

    @property
    def lvl(self):
        return BufferizedNDArrayLevel(self)

    def access(self, indices, op):
        return BufferizedNDArrayAccessor(self).access(indices, op)

    def item(self):
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to Python scalar.")
        val = self.val.load(0)
        return val.item() if hasattr(val, "item") else val

    def __getitem__(self, index):
        """
        Get an item from the bufferized NDArray.
        This allows for indexing into the bufferized array.
        """
        if isinstance(index, slice | np.ndarray) or (
            isinstance(index, tuple) and any(isinstance(i, slice) for i in index)
        ):
            result = self.to_numpy()[index]
            if isinstance(result, np.ndarray):
                return BufferizedNDArray.from_numpy(
                    result, fill_value=self.fill_value, device=self.device
                )
            return Scalar(result, fill_value=self.fill_value, device=self.device)
        if not isinstance(index, tuple) and self.ndim != 1:
            result = self.to_numpy()[index]
            if isinstance(result, np.ndarray):
                return BufferizedNDArray.from_numpy(
                    result, fill_value=self.fill_value, device=self.device
                )
            return Scalar(result, fill_value=self.fill_value, device=self.device)
        if isinstance(index, tuple):
            index = tuple(i for i in index if i is not Ellipsis)
            if index == () and self.ndim == 0:
                return Scalar(
                    self.val.load(0), fill_value=self.fill_value, device=self.device
                )
            if len(index) < self.ndim:
                return BufferizedNDArray.from_numpy(
                    self.to_numpy()[index],
                    fill_value=self.fill_value,
                    device=self.device,
                )
            index = 0 if index == () else np.dot(index, self.strides)
        return Scalar(
            self.val.load(index), fill_value=self.fill_value, device=self.device
        )

    def __setitem__(self, index, value):
        """
        Set an item in the bufferized NDArray.
        This allows for indexing into the bufferized array.
        """
        if isinstance(index, tuple):
            index = tuple(i % s for i, s in zip(index, self._shape, strict=False))
            index = np.ravel_multi_index(index, self._shape)
        self.val.store(index, value)

    def reshape(self, shape, /, *, copy=None):
        return BufferizedNDArray.from_numpy(
            self.to_numpy().reshape(shape),
            fill_value=self.fill_value,
            device=self.device,
        )

    def __str__(self):
        return f"{self.ftype}(shape={tuple(int(d) for d in self.shape)})"

    __repr__ = __str__

element_type property

Data type of the tensor elements.

fill_value property

Default value to fill the tensor.

ftype property

Returns the ftype of the buffer, which is a BufferizedNDArrayFType.

shape_type property

Shape type of the tensor.

__getitem__(index)

Get an item from the bufferized NDArray. This allows for indexing into the bufferized array.

Source code in src/finch/tensor/bufferized_ndarray.py

def __getitem__(self, index):
    """
    Get an item from the bufferized NDArray.
    This allows for indexing into the bufferized array.
    """
    if isinstance(index, slice | np.ndarray) or (
        isinstance(index, tuple) and any(isinstance(i, slice) for i in index)
    ):
        result = self.to_numpy()[index]
        if isinstance(result, np.ndarray):
            return BufferizedNDArray.from_numpy(
                result, fill_value=self.fill_value, device=self.device
            )
        return Scalar(result, fill_value=self.fill_value, device=self.device)
    if not isinstance(index, tuple) and self.ndim != 1:
        result = self.to_numpy()[index]
        if isinstance(result, np.ndarray):
            return BufferizedNDArray.from_numpy(
                result, fill_value=self.fill_value, device=self.device
            )
        return Scalar(result, fill_value=self.fill_value, device=self.device)
    if isinstance(index, tuple):
        index = tuple(i for i in index if i is not Ellipsis)
        if index == () and self.ndim == 0:
            return Scalar(
                self.val.load(0), fill_value=self.fill_value, device=self.device
            )
        if len(index) < self.ndim:
            return BufferizedNDArray.from_numpy(
                self.to_numpy()[index],
                fill_value=self.fill_value,
                device=self.device,
            )
        index = 0 if index == () else np.dot(index, self.strides)
    return Scalar(
        self.val.load(index), fill_value=self.fill_value, device=self.device
    )

__setitem__(index, value)

Set an item in the bufferized NDArray. This allows for indexing into the bufferized array.

Source code in src/finch/tensor/bufferized_ndarray.py

def __setitem__(self, index, value):
    """
    Set an item in the bufferized NDArray.
    This allows for indexing into the bufferized array.
    """
    if isinstance(index, tuple):
        index = tuple(i % s for i, s in zip(index, self._shape, strict=False))
        index = np.ravel_multi_index(index, self._shape)
    self.val.store(index, value)

declare(init, op, shape)

Declare a bufferized NDArray with the given initialization value, operation, and shape.

Source code in src/finch/tensor/bufferized_ndarray.py

def declare(self, init, op, shape):
    """
    Declare a bufferized NDArray with the given initialization value,
    operation, and shape.
    """
    if is_dynamic(self._fill_value):
        # Declaring a tensor sets its dynamic fill to the init.
        init = self.val.ftype.element_type(init)
        self._fill_value = DynamicFill(init, self.val.ftype.element_type)
    elif not bool(ffuncs.same(init, self._fill_value.value)):
        raise ValueError(
            f"Declaring with {init!r} would change the static fill "
            f"{self._fill_value.value!r}"
        )
    for i in range(self.val.length()):
        self.val.store(i, init)
    return self

to_numpy()

Convert the bufferized NDArray to a NumPy array. This is used to get the underlying NumPy array from the bufferized NDArray.

Source code in src/finch/tensor/bufferized_ndarray.py

def to_numpy(self):
    """
    Convert the bufferized NDArray to a NumPy array.
    This is used to get the underlying NumPy array from the bufferized NDArray.
    """
    return self.val.arr.reshape(self._shape, copy=False)

BufferizedNDArrayFType

Bases: FinchTensorFType, MutableStructFType

A ftype for bufferized NumPy arrays that provides metadata about the array. This includes the fill value, element type, and shape type.

Source code in src/finch/tensor/bufferized_ndarray.py

class BufferizedNDArrayFType(FinchTensorFType, MutableStructFType):
    """
    A ftype for bufferized NumPy arrays that provides metadata about the array.
    This includes the fill value, element type, and shape type.
    """

    @property
    def struct_name(self):
        return "BufferizedNDArray"

    @property
    def struct_fields(self):
        fields = [
            ("val", self.buf_t),
            ("shape", self.shape_t),
            ("strides", self.strides_t),
        ]
        if is_dynamic(self._fill_type):
            # The fill value is bound at call time through a struct field.
            fields.append(("fill", self.element_type))
        return fields

    def struct_getattr(self, obj, attr):
        # A dynamic fill is bound at call time through the "fill" field.
        if attr == "fill":
            return obj.fill_value
        return getattr(obj, attr)

    def struct_setattr(self, obj, attr, value):
        match attr:
            case "shape":
                obj._shape = value
            case "fill":
                obj._fill_value = DynamicFill(value, self.element_type)
            case _:
                setattr(obj, attr, value)

    @property
    def lvl_t(self):
        return BufferizedNDArrayLevelFType(self)

    def get_child_type(self, attr):
        if attr != "lvl":
            raise TypeError(f"{self} does not support child {attr!r}")
        return self.lvl_t

    def get_child(self, obj, attr):
        self.get_child_type(attr)
        return obj

    def from_fields(self, buf, shape, strides, fill=None):
        if is_dynamic(self.fill_type):
            # The fill field holds the value of the dynamic fill.
            fill = DynamicFill(fill, self.element_type)
        return BufferizedNDArray(
            buf, shape, strides, resolve_fill(self.fill_type, fill), device=self.device
        )

    def from_numpy(self, arr: np.ndarray, fill_value=None):
        val = NumpyBuffer(arr.reshape(-1, copy=False))
        strides = _get_default_strides(arr.shape)
        return BufferizedNDArray(
            val=val,
            shape=tuple(
                t(s)
                for s, t in zip(arr.shape, self.shape_t.struct_fieldtypes, strict=True)
            ),
            strides=tuple(
                t(s)
                for (s, t) in zip(
                    strides, self.strides_t.struct_fieldtypes, strict=True
                )
            ),
            fill_value=resolve_fill(
                self.fill_type, None if fill_value is None else as_fill(fill_value)
            ),
            device=self.device,
        )

    def __init__(
        self,
        *,
        buffer_type: NumpyBufferFType,
        ndim: int,
        dimension_type: TupleFType | tuple[FType, ...],
        fill_type: AbstractFillFType | None = None,
        device=None,
    ):
        if not isinstance(dimension_type, TupleFType):
            dimension_type = TupleFType.from_tuple(dimension_type)
        assert isinstance(dimension_type, TupleFType)
        # Normalize dimension field types to Finch ftypes so generated
        # result_type values are consistent with strict asm type checks.
        dimension_type = TupleFType.from_tuple(
            tuple(ftype(t) for t in dimension_type.struct_fieldtypes)
        )
        self.buf_t = buffer_type
        self._ndim = ndim
        self.shape_t = dimension_type
        self.strides_t = TupleFType.from_tuple((self.buf_t.length_type,) * ndim)
        elem_t = self.buf_t.element_type
        # The fill defaults to a static zero, and takes the buffer's element type.
        if fill_type is None:
            fill_type = StaticFillFType(elem_t(0))
        self._fill_type = fill_type.astype(elem_t)
        self._device = normalize_device(device)

    def construct(
        self,
        shape: tuple[int, ...],
        fill_value: Any = None,
    ) -> BufferizedNDArray:
        # The fill value may be raw, which is a static fill.
        fill = resolve_fill(
            self.fill_type, None if fill_value is None else as_fill(fill_value)
        )
        arr = np.empty(shape, dtype=np_dtype(self.element_type))
        arr[...] = fill.value
        return self.from_numpy(arr, fill_value=fill)

    def __call__(
        self,
        val: Any,
    ) -> BufferizedNDArray:
        """
        Convert a tensor to this bufferized ndarray type.

        Args:
            val: A tensor to convert to this type.
        Returns:
            A BufferizedNDArray instance of this type.
        """
        raise NotImplementedError(
            f"Tensor conversion not yet implemented for {type(self).__name__}"
        )

    def __eq__(self, other):
        if not isinstance(other, BufferizedNDArrayFType):
            return False
        return (
            self.buf_t == other.buf_t
            and self.ndim == other.ndim
            and self.shape_t == other.shape_t
            and self.fill_type == other.fill_type
            and self.device == other.device
        )

    def with_fill(self, fill_type: AbstractFillFType) -> "BufferizedNDArrayFType":
        return BufferizedNDArrayFType(
            buffer_type=self.buf_t,
            ndim=self.ndim,
            dimension_type=self.shape_t,
            fill_type=fill_type,
            device=self.device,
        )

    def fisinstance(self, other):
        other_t = ftype(other)
        if is_dynamic(self._fill_type) and isinstance(other_t, BufferizedNDArrayFType):
            # A dynamic-fill ftype accepts any fill value of matching dtype.
            other_t = other_t.with_fill(self._fill_type)
        return other_t == self

    def __hash__(self):
        return hash(
            (
                self.buf_t,
                self.ndim,
                self.shape_t,
                self.fill_type,
                self.device,
            )
        )

    def __str__(self):
        return str(self.struct_name)

    def __repr__(self):
        return (
            f"BufferizedNDArrayFType(buffer_type={repr(self.buf_t)},"
            f" ndim = {self.ndim}, dimension_type ={repr(self.shape_t)},"
            f" fill_type={self.fill_type!r})"
        )

    @property
    def ndim(self) -> int:
        return int(self._ndim)

    @property
    def fill_type(self) -> AbstractFillFType:
        return self._fill_type

    @property
    def device(self):
        return self._device

    @property
    def element_type(self):
        return self.buf_t.element_type

    @property
    def shape_type(self) -> tuple:
        return tuple(self.shape_t.struct_fieldtypes)

    @property
    def level_format_properties(self) -> list[FormatProperty]:
        return [Dense(tuple(range(n + 1))) for n in range(self.ndim)]

    def lower_dim(self, ctx, obj, r):
        return self.lvl_t.level_lower_dim(ctx, ctx(obj.lvl), r)

    def lower_fill(self, ctx, obj):
        return self.lvl_t.lower_fill(ctx(obj.lvl))

    def lower_declare(self, ctx, tns: ntn.Fiber, init, op, shape):
        root = ctx(tns.lvl.root)
        if is_dynamic(self.fill_type):
            # The init becomes the fill of the tensor.
            if init.result_type != self.element_type:
                init = asm.Call(asm.Literal(ffuncs.astype(self.element_type)), (init,))
            ctx.exec(asm.SetAttr(root, asm.Literal("fill"), init))
            init = asm.GetAttr(root, asm.Literal("fill"))
        i_var = asm.Variable("i", self.buf_t.length_type)
        buf = asm.GetAttr(root, asm.Literal("val"))
        body = asm.Store(
            buf,
            i_var,
            init,
        )
        ctx.exec(asm.ForLoop(i_var, asm.Literal(np.intp(0)), asm.Length(buf), body))
        return

    def lower_freeze(self, ctx, tns, op):
        return tns

    def lower_thaw(self, ctx, tns, op):
        return tns

    def unfurl(self, ctx, tns, ext, mode, proto):
        tns = ctx.resolve(tns)
        return self.lvl_t.level_unfurl(ctx, tns, ext, mode, proto, ctx(tns.pos))

    def reshape(self, arr: BufferizedNDArray, new_shape: tuple[np.intp, ...]):
        old_size = functools.reduce(operator.mul, arr.shape, np.intp(1))
        new_size = functools.reduce(operator.mul, new_shape, np.intp(1))
        if old_size != new_size:
            raise ValueError(
                f"Cannot reshape array of size {old_size} into shape {new_shape}"
            )
        new_strides = _get_default_strides(new_shape)
        default_strides = _get_default_strides(arr.shape)
        if arr.strides == default_strides:
            return BufferizedNDArray(
                arr.val,
                new_shape,
                new_strides,
                fill_value=arr.fill_value,
                device=arr.device,
            )
        return BufferizedNDArray.from_numpy(
            arr.to_numpy().reshape(new_shape),
            fill_value=arr.fill_value,
            device=arr.device,
        )

    def lower_unwrap(self, ctx, obj):
        return self.lvl_t.level_lower_unwrap(ctx, obj, ctx(obj.pos))

    def lower_increment(self, ctx, obj, op, val):
        return self.lvl_t.level_lower_increment(ctx, obj, op, val, ctx(obj.pos))

__call__(val)

Convert a tensor to this bufferized ndarray type.

Parameters:

Name	Type	Description	Default
val	Any	A tensor to convert to this type.	required

Returns: A BufferizedNDArray instance of this type.

Source code in src/finch/tensor/bufferized_ndarray.py

def __call__(
    self,
    val: Any,
) -> BufferizedNDArray:
    """
    Convert a tensor to this bufferized ndarray type.

    Args:
        val: A tensor to convert to this type.
    Returns:
        A BufferizedNDArray instance of this type.
    """
    raise NotImplementedError(
        f"Tensor conversion not yet implemented for {type(self).__name__}"
    )

ChunkMaskTensor

Bases: PatternTensor

Map n indices to chunks of size b.

shape is (ceil(n / b), n), and row j holds the indices of chunk j, so that the efficient row-major traversal visits each chunk and then its indices. The last chunk may be shorter.

Source code in src/finch/tensor/patterns.py

class ChunkMaskTensor(PatternTensor):
    """Map ``n`` indices to chunks of size ``b``.

    ``shape`` is ``(ceil(n / b), n)``, and row ``j`` holds the indices of chunk
    ``j``, so that the efficient row-major traversal visits each chunk and then
    its indices. The last chunk may be shorter.
    """

    def __init__(self, shape, *, b: int, dtype=None):
        self._b = operator.index(b)
        if self._b <= 0:
            raise ValueError("b must be positive")
        super().__init__(
            shape,
            dtype=dtype,
            default_dtype=np.bool_,
            fill_value=False,
            pattern_value=True,
            b=self._b,
        )
        chunks, n = map(operator.index, self.shape)
        if n < 0:
            raise ValueError("n must be nonnegative")
        if chunks != (n + self._b - 1) // self._b:
            raise ValueError("shape[0] must equal ceil(shape[1] / b)")

    def contains(self, j, i) -> bool:
        return j == i // self._b

ConstantScalar

Bases: Scalar

A scalar whose value is treated as a compile-time constant. Kernels may specialize on the value (and recompile per distinct value), unlike plain Scalars, whose values are bound at call time.

Source code in src/finch/tensor/scalar.py

class ConstantScalar(Scalar):
    """
    A scalar whose value is treated as a compile-time constant. Kernels may
    specialize on the value (and recompile per distinct value), unlike plain
    `Scalar`s, whose values are bound at call time.
    """

    def __init__(self, val: Any, device=None):
        super().__init__(val, device=device)

    @property
    def ftype(self):
        return ScalarFType(
            ftype(self.val), StaticFillFType(self._fill_value.value), self._device
        )

DC dataclass

A degree constraint: value bounds ||deg(to_indices | from_indices)||_p.

p = inf (the default) is the classic max-degree constraint; p = 1 is the distinct count; intermediate p capture skew between the two.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

@dataclass(frozen=True)
class DC:
    """
    A degree constraint: ``value`` bounds ``||deg(to_indices | from_indices)||_p``.

    ``p = inf`` (the default) is the classic max-degree constraint; ``p = 1`` is
    the distinct count; intermediate ``p`` capture skew between the two.
    """

    from_indices: frozenset[Field]
    to_indices: frozenset[Field]
    value: float
    p: float = math.inf

DCStats

Bases: BoundStats

Classic degree constraints -- max degrees (p = inf) and distinct counts. Estimates size with a best-first product search over the constraints, clamped by the tensor's dense capacity.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

class DCStats(BoundStats):
    """
    Classic degree constraints -- max degrees (``p = inf``) and distinct counts.
    Estimates size with a best-first product search over the constraints,
    clamped by the tensor's dense capacity.
    """

    default_ps: ClassVar[tuple[float, ...]] = (math.inf,)

    def estimate_non_fill_values(self) -> float:
        """Smallest degree-constraint product covering all indices, clamped by
        dense capacity."""
        idx: frozenset[Field] = frozenset(self.dim_sizes.keys())
        if len(idx) == 0:
            return 1.0

        best: dict[frozenset[Field], float] = {frozenset(): 1.0}
        frontier: set[frozenset[Field]] = {frozenset()}

        while True:
            current_bound = best.get(idx, math.inf)
            new_frontier: set[frozenset[Field]] = set()

            for node in frontier:
                for dc in self.dcs:
                    if node.issuperset(dc.from_indices):
                        y = node.union(dc.to_indices)
                        if best[node] > float(2 ** (64 - 2)) or float(dc.value) > float(
                            2 ** (64 - 2)
                        ):
                            y_weight = float(2**64)
                        else:
                            y_weight = best[node] * dc.value
                        if min(current_bound, best.get(y, math.inf)) > y_weight:
                            best[y] = y_weight
                            new_frontier.add(y)
            if len(new_frontier) == 0:
                break
            frontier = new_frontier

        min_weight = float(self.get_dim_space_size(idx))
        for node, weight in best.items():
            if node.issuperset(idx):
                min_weight = min(min_weight, weight)
        return min_weight

estimate_non_fill_values()

Smallest degree-constraint product covering all indices, clamped by dense capacity.

Source code in src/finch/autoschedule/tensor_stats/bound_stats.py

def estimate_non_fill_values(self) -> float:
    """Smallest degree-constraint product covering all indices, clamped by
    dense capacity."""
    idx: frozenset[Field] = frozenset(self.dim_sizes.keys())
    if len(idx) == 0:
        return 1.0

    best: dict[frozenset[Field], float] = {frozenset(): 1.0}
    frontier: set[frozenset[Field]] = {frozenset()}

    while True:
        current_bound = best.get(idx, math.inf)
        new_frontier: set[frozenset[Field]] = set()

        for node in frontier:
            for dc in self.dcs:
                if node.issuperset(dc.from_indices):
                    y = node.union(dc.to_indices)
                    if best[node] > float(2 ** (64 - 2)) or float(dc.value) > float(
                        2 ** (64 - 2)
                    ):
                        y_weight = float(2**64)
                    else:
                        y_weight = best[node] * dc.value
                    if min(current_bound, best.get(y, math.inf)) > y_weight:
                        best[y] = y_weight
                        new_frontier.add(y)
        if len(new_frontier) == 0:
            break
        frontier = new_frontier

    min_weight = float(self.get_dim_space_size(idx))
    for node, weight in best.items():
        if node.issuperset(idx):
            min_weight = min(min_weight, weight)
    return min_weight

DenseLevel dataclass

Bases: SingleDimensionLevel

A class representing dense level.

Source code in src/finch/tensor/level/dense_level.py

@dataclass
class DenseLevel(SingleDimensionLevel):
    """
    A class representing dense level.
    """

    lvl: Level
    dimension: np.integer
    stride: np.integer = field(init=False)

    def __post_init__(self):
        # The stride is stored so that kernels can write it back; it must be
        # recomputed whenever the dimensions of the dense levels below change.
        dim_t = ftype(self.dimension)
        match self.lvl:
            case DenseLevel():
                self.stride = dim_t(self.lvl.dimension * self.lvl.stride)
            case _:
                self.stride = dim_t(1)

    @property
    def shape(self) -> tuple:
        return (self.dimension, *self.lvl.shape)

    @property
    def ftype(self) -> DenseLevelFType:
        return DenseLevelFType(self.lvl.ftype, ftype(self.dimension))

    @property
    def val(self) -> Any:
        return self.lvl.val

    def __str__(self):
        return f"DenseLevel(lvl={self.lvl}, dim={self.dimension})"

DenseLevelFType dataclass

Bases: SingleDimensionLevelFType, MutableStructFType

Source code in src/finch/tensor/level/dense_level.py

@dataclass(unsafe_hash=True)
class DenseLevelFType(SingleDimensionLevelFType, MutableStructFType):
    _lvl_t: LevelFType
    dimension_type: ftypes.FDTypeInteger = ftypes.intp

    @property
    def struct_name(self):
        return "DenseLevelFType"

    @property
    def struct_fields(self):
        return [
            ("lvl", self.lvl_t),
            ("dimension", self.dimension_type),
            ("stride", self.dimension_type),
        ]

    def struct_setattr(self, obj, attr, value):
        setattr(obj, attr, value)

    def level_cost(self, fields, stats, stats_factory, num_pos, lvl) -> float:
        n = stats.get_dim_size(fields[lvl])
        return self.lvl_t.level_cost(fields, stats, stats_factory, num_pos * n, lvl + 1)

    def __post_init__(self):
        self.dimension_type = ftype(self.dimension_type)

    def construct(self, shape: tuple[Any, ...], *, pos: int) -> DenseLevel:
        """
        Creates an instance of DenseLevel with the given ftype.

        Args:
            shape: The shape to be used for the level. (mandatory)
        Returns:
            An instance of DenseLevel.
        """
        dimension = self.dimension_type(shape[0])
        lvl = self.lvl_t.construct(
            shape=shape[1:],
            pos=int(pos) * int(dimension),
        )
        return DenseLevel(lvl, dimension)

    def __call__(self, val: Any) -> DenseLevel:
        """
        Convert a level to this dense level type.

        Args:
            val: A value to convert to this type.
        Returns:
            A DenseLevel instance of this type.
        """
        raise NotImplementedError(
            f"Level conversion not yet implemented for {type(self).__name__}"
        )

    def __str__(self):
        return f"DenseLevelFType({self.lvl_t})"

    @property
    def ndim(self):
        return 1 + self.lvl_t.ndim

    @property
    def fill_type(self):
        return self.lvl_t.fill_type

    @property
    def element_type(self):
        """
        Returns the type of elements stored in the fibers.
        """
        return self.lvl_t.element_type

    @property
    def shape_type(self):
        """
        Returns the type of the shape of the fibers.
        """
        return (self.dimension_type, *self.lvl_t.shape_type)

    @property
    def position_type(self):
        """
        Returns the type of positions within the levels.
        """
        return self.lvl_t.position_type

    @property
    def buffer_type(self):
        return self.lvl_t.buffer_type

    @property
    def buffer_factory(self):
        """
        Returns the ftype of the buffer used for the fibers.
        """
        return self.lvl_t.buffer_factory

    @property
    def lvl_t(self):
        return self._lvl_t

    def level_format_properties(self, n):
        return [Dense(tuple(range(n + 1)))] + self.lvl_t.level_format_properties(n + 1)

    def level_lower_dim(self, ctx, obj, r):
        if r == 0:
            return asm.GetAttr(obj, asm.Literal("dimension"))
        return self.lvl_t.level_lower_dim(
            ctx, asm.GetAttr(obj, asm.Literal("lvl")), r - 1
        )

    def level_lower_declare(self, ctx, tns, init, op, shape, pos):
        pos = asm.Call(
            asm.Literal(ffuncs.mul),
            (pos, asm.GetAttr(tns, asm.Literal("dimension"))),
        )
        return self.lvl_t.level_lower_declare(
            ctx, asm.GetAttr(tns, asm.Literal("lvl")), init, op, shape, pos
        )

    def level_lower_freeze(self, ctx, lvl, op, pos):
        pos = asm.Call(
            asm.Literal(ffuncs.mul),
            (pos, asm.GetAttr(lvl, asm.Literal("dimension"))),
        )
        return self.lvl_t.level_lower_freeze(
            ctx, asm.GetAttr(lvl, asm.Literal("lvl")), op, pos
        )

    def level_lower_thaw(self, ctx, lvl, op, pos):
        pos = asm.Call(
            asm.Literal(ffuncs.mul),
            (pos, asm.GetAttr(lvl, asm.Literal("dimension"))),
        )
        return self.lvl_t.level_lower_thaw(
            ctx, asm.GetAttr(lvl, asm.Literal("lvl")), op, pos
        )

    def level_lower_assemble(self, ctx, lvl, start, stop):
        dimension = asm.GetAttr(lvl, asm.Literal("dimension"))
        self.lvl_t.level_lower_assemble(
            ctx,
            asm.GetAttr(lvl, asm.Literal("lvl")),
            asm.Call(asm.Literal(ffuncs.mul), (start, dimension)),
            asm.Call(asm.Literal(ffuncs.mul), (stop, dimension)),
        )

    def level_lower_increment(self, ctx, obj, op, val, pos):
        raise NotImplementedError(
            "DenseLevelFType does not support level_lower_increment."
        )

    def level_lower_unwrap(self, ctx, obj, pos):
        raise NotImplementedError(
            "DenseLevelFType does not support level_lower_unwrap."
        )

    def level_unfurl(
        self,
        ctx: AssemblyContext,
        fiber: ntn.Fiber,
        ext: SymbolicExtent,
        mode,
        proto,
        pos: asm.AssemblyExpression,
    ):
        tns = fiber
        level = tns.lvl
        lvl = ctx(level)

        def child_accessor(ctx: LoopletContext, idx: ntn.Variable):
            if idx.type_ is None:
                raise TypeError(f"Expected loop variable type for {idx.name}")
            pos_2 = asm.Variable(
                ctx.freshen(idx, f"_pos_{self.ndim - 1}"), self.position_type
            )
            child_pos = asm.Call(
                asm.Literal(ffuncs.add),
                (
                    asm.Call(
                        asm.Literal(ffuncs.mul),
                        (pos, asm.GetAttr(lvl, asm.Literal("dimension"))),
                    ),
                    ctx.ctx(idx),
                ),
            )
            if child_pos.result_type != self.position_type:
                child_pos = asm.Call(
                    asm.Literal(ffuncs.astype(self.position_type)), (child_pos,)
                )
            ctx.exec(asm.Assign(pos_2, child_pos))
            return lplt.Run(
                replace(
                    tns,
                    lvl=ntn.Child(level),
                    pos=ntn.Value(pos_2, self.position_type),
                    idxs=(*tns.idxs, idx),
                )
            )

        return lplt.Lookup(child_accessor)

    def from_fields(self, lvl, dimension, stride) -> DenseLevel:
        return DenseLevel(lvl=lvl, dimension=dimension)

buffer_factory property

Returns the ftype of the buffer used for the fibers.

element_type property

Returns the type of elements stored in the fibers.

position_type property

Returns the type of positions within the levels.

shape_type property

Returns the type of the shape of the fibers.

__call__(val)

Convert a level to this dense level type.

Parameters:

Name	Type	Description	Default
val	Any	A value to convert to this type.	required

Returns: A DenseLevel instance of this type.

Source code in src/finch/tensor/level/dense_level.py

def __call__(self, val: Any) -> DenseLevel:
    """
    Convert a level to this dense level type.

    Args:
        val: A value to convert to this type.
    Returns:
        A DenseLevel instance of this type.
    """
    raise NotImplementedError(
        f"Level conversion not yet implemented for {type(self).__name__}"
    )

construct(shape, *, pos)

Creates an instance of DenseLevel with the given ftype.

Parameters:

Name	Type	Description	Default
shape	tuple[Any, ...]	The shape to be used for the level. (mandatory)	required

Returns: An instance of DenseLevel.

Source code in src/finch/tensor/level/dense_level.py

def construct(self, shape: tuple[Any, ...], *, pos: int) -> DenseLevel:
    """
    Creates an instance of DenseLevel with the given ftype.

    Args:
        shape: The shape to be used for the level. (mandatory)
    Returns:
        An instance of DenseLevel.
    """
    dimension = self.dimension_type(shape[0])
    lvl = self.lvl_t.construct(
        shape=shape[1:],
        pos=int(pos) * int(dimension),
    )
    return DenseLevel(lvl, dimension)

ElementLevel dataclass

Bases: Level

A class representing the leaf level of Finch tensors.

Source code in src/finch/tensor/level/element_level.py

@dataclass
class ElementLevel(Level):
    """
    A class representing the leaf level of Finch tensors.
    """

    _format: ElementLevelFType = field(repr=False)
    _val: Any | None = None
    # The fill belongs to the level, so that a dynamic fill may change without
    # changing the format. Only a static fill can default to that of the format.
    _fill: Any = None

    def __post_init__(self):
        if self._val is None:
            self._val = self._format.buffer_type(len=0)
        if self._fill is None:
            fill_type = self._format.fill_type
            assert fill_type is not None
            self._fill = resolve_fill(fill_type).value

    @property
    def shape(self) -> tuple:
        return ()

    @property
    def ftype(self) -> ElementLevelFType:
        return self._format

    def with_fill(self, fill_value: AbstractFill) -> "ElementLevel":
        return replace(
            self,
            _format=self._format.with_fill(fill_value.ftype),
            _fill=fill_value.value,
        )

    @property
    def val(self) -> Any:
        return self._val

    @property
    def fill_value(self) -> Any:
        return self._fill

    def __str__(self):
        return f"ElementLevel(val={self._val})"

ElementLevelFType dataclass

Bases: LevelFType, MutableStructFType

Source code in src/finch/tensor/level/element_level.py

@dataclass(unsafe_hash=True)
class ElementLevelFType(LevelFType, MutableStructFType):
    fill_type: AbstractFillFType | None = None
    element_type: FType | None = None
    position_type: FType | None = None
    buffer_factory: Any = NumpyBufferFType
    buffer_type: Any = None

    @property
    def struct_name(self):
        return "ElementLevelFType"

    @property
    def struct_fields(self):
        fields = [
            ("val", self.buffer_type),
        ]
        if is_dynamic(self.fill_type):
            # The fill value is bound at call time through a struct field.
            fields.append(("fill", self.element_type))
        return fields

    def struct_getattr(self, obj, attr):
        # A dynamic fill is bound at call time through the "fill" field.
        if attr == "fill":
            return obj.fill_value
        return getattr(obj, attr)

    def struct_setattr(self, obj, attr, value):
        match attr:
            case "val":
                obj._val = value
            case "fill":
                obj._fill = value

    def level_cost(self, fields, stats, stats_factory, num_pos, lvl) -> float:
        # no inner level
        # cost = num_pos * bytes per value
        assert self.element_type is not None
        val_size = np_dtype(self.element_type).itemsize
        return num_pos * val_size

    def __post_init__(self):
        # Ensure element_type is an FType
        if self.element_type is None:
            assert self.fill_type is not None, (
                "Must provide either element_type or fill_type."
            )
            self.element_type = self.fill_type.element_type
        assert isinstance(self.element_type, FType), (
            "element_type must be an instance of FType"
        )
        if self.buffer_type is None:
            self.buffer_type = self.buffer_factory(self.element_type)
        if self.position_type is None:
            self.position_type = ftype(np.intp)
        self.position_type = ftype(self.position_type)
        self.element_type = self.buffer_type.element_type
        # The fill defaults to a static zero, and takes the buffer's element type.
        fill_type = self.fill_type
        if fill_type is None:
            fill_type = StaticFillFType(self.element_type(0))
        self.fill_type = fill_type.astype(self.element_type)

    def construct(self, shape: tuple[Any, ...], *, pos: int) -> "ElementLevel":
        """
        Creates an instance of ElementLevel with the given ftype.

        Args:
            shape: Should be always `()`, used for validation.
        Returns:
            An instance of ElementLevel.
        """
        if len(shape) != 0:
            raise ValueError("ElementLevelFType must be called with an empty shape.")
        val = self.buffer_type(len=pos)
        return ElementLevel(self, val)

    def __call__(self, val: Any) -> "ElementLevel":
        """
        Convert a level to this element level type.

        Args:
            val: A value to convert to this type.
        Returns:
            An ElementLevel instance of this type.
        """
        raise NotImplementedError(
            f"Level conversion not yet implemented for {type(self).__name__}"
        )

    def __str__(self):
        match self.fill_type:
            case StaticFillFType() as fill_type:
                return f"ElementLevelFType(fv={fill_type.value})"
            case fill_type:
                return f"ElementLevelFType(fv={fill_type})"

    @property
    def ndim(self):
        return 0

    @property
    def shape_type(self):
        return ()

    @property
    def lvl_t(self):
        raise Exception("ElementLevelFType is the leaf level.")

    def level_format_properties(self, n):
        return []

    def with_fill(self, fill_type: Any) -> "ElementLevelFType":
        return ElementLevelFType(
            fill_type=fill_type,
            element_type=self.element_type,
            position_type=self.position_type,
            buffer_factory=self.buffer_factory,
            buffer_type=self.buffer_type,
        )

    def lower_fill(self, lvl_expr):
        return asm.GetAttr(lvl_expr, asm.Literal("fill"))

    def from_fields(self, val=None, fill=None) -> "ElementLevel":
        # Wrap numpy arrays in NumpyBuffer and flatten, similar to BufferizedNDArray
        if val is not None and isinstance(val, np.ndarray):
            from finch.codegen import NumpyBuffer

            val = NumpyBuffer(np.asarray(val).reshape(-1, copy=False))
        return ElementLevel(_format=self, _val=val, _fill=fill)

    def level_lower_declare(self, ctx, tns, init, op, shape, pos):
        if is_dynamic(self.fill_type):
            # The init becomes the fill of the tensor.
            elem_t = self.fill_type.element_type
            if init.result_type != elem_t:
                init = asm.Call(asm.Literal(ffuncs.astype(elem_t)), (init,))
            ctx.exec(asm.SetAttr(tns, asm.Literal("fill"), init))
            init = asm.GetAttr(tns, asm.Literal("fill"))
        buf = asm.GetAttr(tns, asm.Literal("val"))
        i_var = asm.Variable("i", self.buffer_type.length_type)
        body = asm.Store(buf, i_var, init)
        ctx.exec(asm.ForLoop(i_var, asm.Literal(np.intp(0)), asm.Length(buf), body))

    def level_lower_unwrap(self, ctx, obj, pos):
        buf = asm.GetAttr(ctx(obj.lvl), asm.Literal("val"))
        if pos.result_type != self.buffer_type.length_type:
            pos = asm.Call(
                asm.Literal(ffuncs.astype(self.buffer_type.length_type)), (pos,)
            )
        return asm.Load(buf, pos)

    def level_lower_increment(
        self,
        ctx: AssemblyContext,
        obj,
        op: ntn.Literal,
        val: ntn.NotationExpression,
        pos: asm.AssemblyExpression,
    ):
        buf = asm.GetAttr(ctx(obj.lvl), asm.Literal("val"))
        if pos.result_type != self.buffer_type.length_type:
            pos = asm.Call(
                asm.Literal(ffuncs.astype(self.buffer_type.length_type)), (pos,)
            )
        pos_e, op_e, val_e = pos, ctx(op), ctx(val)
        ctx.exec(
            asm.Store(
                buf,
                pos_e,
                asm.Call(op_e, (asm.Load(buf, pos_e), val_e)),
            )
        )
        match obj:
            case ntn.HollowFiber(dirty=dirty):
                fill = (
                    asm.Literal(self.fill_type.value)
                    if is_static(self.fill_type)
                    else self.lower_fill(ctx(obj.lvl))
                )
                ctx.exec(
                    asm.If(
                        asm.Call(
                            asm.Literal(ffuncs.not_),
                            (
                                asm.Call(
                                    asm.Literal(ffuncs.same),
                                    (asm.Load(buf, pos_e), fill),
                                ),
                            ),
                        ),
                        asm.Assign(ctx(dirty), asm.Literal(True)),
                    )
                )

    def level_lower_assemble(self, ctx, lvl, start, stop):
        buf = asm.GetAttr(lvl, asm.Literal("val"))
        p_t = self.buffer_type.length_type
        to_size = asm.Literal(ffuncs.astype(p_t))
        p = asm.Variable(ctx.freshen("p"), p_t)
        p_start = asm.Variable(ctx.freshen("p_start"), p_t)
        p_stop = asm.Variable(ctx.freshen("p_stop"), p_t)
        length = asm.Length(buf)
        fill = (
            asm.Literal(self.fill_type.value)
            if is_static(self.fill_type)
            else self.lower_fill(lvl)
        )
        expr = """finch
        p_start = to_size(start)
        p_stop = to_size(stop)
        if (length < p_stop)
            resize(buf, p_stop)
        end
        for (p in p_start:p_stop)
            buf[p] = fill
        end
        """
        ctx.exec(parse_assembly(expr, locals(), position_type=p_t))

    def level_lower_freeze(self, ctx, lvl, op, pos):
        buf = asm.GetAttr(lvl, asm.Literal("val"))
        size = asm.Call(
            asm.Literal(ffuncs.astype(self.buffer_type.length_type)), (pos,)
        )
        length = asm.Length(buf)
        expr = """finch
        if (length != size)
            resize(buf, size)
        end
        """
        ctx.exec(parse_assembly(expr, locals()))

    def level_lower_thaw(self, ctx, lvl, op, pos):
        pass

    def level_lower_dim(self, ctx, obj, r):
        raise NotImplementedError("ElementLevelFType does not support level_lower_dim.")

    def level_unfurl(self, ctx, lvl, ext, mode, proto, pos):
        raise NotImplementedError("ElementLevelFType does not support level_unfurl.")

__call__(val)

Convert a level to this element level type.

Parameters:

Name	Type	Description	Default
val	Any	A value to convert to this type.	required

Returns: An ElementLevel instance of this type.

Source code in src/finch/tensor/level/element_level.py

def __call__(self, val: Any) -> "ElementLevel":
    """
    Convert a level to this element level type.

    Args:
        val: A value to convert to this type.
    Returns:
        An ElementLevel instance of this type.
    """
    raise NotImplementedError(
        f"Level conversion not yet implemented for {type(self).__name__}"
    )

construct(shape, *, pos)

Creates an instance of ElementLevel with the given ftype.

Parameters:

Name	Type	Description	Default
shape	tuple[Any, ...]	Should be always (), used for validation.	required

Returns: An instance of ElementLevel.

Source code in src/finch/tensor/level/element_level.py

def construct(self, shape: tuple[Any, ...], *, pos: int) -> "ElementLevel":
    """
    Creates an instance of ElementLevel with the given ftype.

    Args:
        shape: Should be always `()`, used for validation.
    Returns:
        An instance of ElementLevel.
    """
    if len(shape) != 0:
        raise ValueError("ElementLevelFType must be called with an empty shape.")
    val = self.buffer_type(len=pos)
    return ElementLevel(self, val)

Extent dataclass

Bases: FTyped[ExtentFType]

A class to represent the extent of a loop variable. This is used to define the start and end values of a loop.

The extent is start-inclusive and end-exclusive, meaning that the loop variable will take on values from start to end-1.

Source code in src/finch/compile/lower.py

@dataclass(frozen=True)
class Extent(FTyped[ExtentFType]):
    """
    A class to represent the extent of a loop variable. This is used to define
    the start and end values of a loop.

    The extent is start-inclusive and end-exclusive, meaning that the loop variable
    will take on values from start to end-1.
    """

    start: Any
    end: Any

    def loop_interpreter(self, ctx, idx, body):
        for idx_e in range(self.start, self.end):
            # Create a new scope for each iteration
            ctx_2 = ctx.scope(loop_state=HaltState())
            # Assign the loop variable
            ctx_2.bindings[idx.name] = idx.type_(idx_e)
            # Execute the body of the loop
            ctx_2(body)

    @property
    def ftype(self):
        return ExtentFType(ftype(self.start), ftype(self.end))

ExtentFType dataclass

Bases: ImmutableStructFType

Source code in src/finch/compile/lower.py

@dataclass(eq=True, frozen=True)
class ExtentFType(ImmutableStructFType):
    start_t: FType
    end_t: FType

    def __repr__(self):
        return f"ExtentFType(start={self.start_t}, end={self.end_t})"

    @property
    def struct_name(self):
        return "Extent"

    @property
    def struct_fields(self):
        return [("start", self.start_t), ("end", self.end_t)]

    def from_fields(self, start, end) -> Extent:
        return Extent(start, end)

    def __call__(self, *args):
        raise TypeError(f"{self.struct_name} is not callable")

    def lower_loop(
        self,
        ctx: AssemblyContext,
        idx: ntn.Variable,
        ext: SymbolicExtent,
        body: ntn.NotationExpression,
    ):
        """
        Lower a loop with the given index and body.
        This is used to compile the loop into assembly.
        """
        lower_looplets(ctx, idx, ext, body)
        return

    def default_loop(self, ctx, idx, ext: SymbolicExtent, body):
        def assert_lowered(node):
            match node:
                case ntn.Access(_, _, (j, *_)):
                    if j == idx:
                        raise FinchCompileError(
                            node, f"Access with {j} should have been lowered already"
                        )
            return

        for node in PostOrderDFS(body):
            assert_lowered(node)

        ctx_2 = ctx.scope()
        ctx_2(body)
        body_3 = asm.Block(ctx_2.emit())
        ctx.exec(
            asm.ForLoop(
                ctx(idx),
                ctx(ext.get_start()),
                ctx(ext.get_end()),
                body_3,
            )
        )
        return

lower_loop(ctx, idx, ext, body)

Lower a loop with the given index and body. This is used to compile the loop into assembly.

Source code in src/finch/compile/lower.py

def lower_loop(
    self,
    ctx: AssemblyContext,
    idx: ntn.Variable,
    ext: SymbolicExtent,
    body: ntn.NotationExpression,
):
    """
    Lower a loop with the given index and body.
    This is used to compile the loop into assembly.
    """
    lower_looplets(ctx, idx, ext, body)
    return

FType

Bases: ABC

Source code in src/finch/algebra/ftypes.py

class FType(ABC):
    @abstractmethod
    def __eq__(self, other) -> builtins.bool: ...

    @abstractmethod
    def __hash__(self) -> builtins.int: ...

    @abstractmethod
    def __call__(self, val: Any) -> Any:
        """
        Perform type conversion or construction with the given value.

        For numeric/scalar types, this converts the value to this type.
        For composite types, this may construct an instance of this type.

        Args:
            val: The value to convert or use for construction

        Returns:
            The converted/constructed value
        """
        ...

    def fisinstance(self, other) -> builtins.bool:
        """
        Check if `other` is an instance of this ftype.
        """
        return ftype(other) == self

__call__(val) abstractmethod

Perform type conversion or construction with the given value.

For numeric/scalar types, this converts the value to this type. For composite types, this may construct an instance of this type.

Parameters:

Name	Type	Description	Default
val	Any	The value to convert or use for construction	required

Returns:

Type	Description
Any	The converted/constructed value

Source code in src/finch/algebra/ftypes.py

@abstractmethod
def __call__(self, val: Any) -> Any:
    """
    Perform type conversion or construction with the given value.

    For numeric/scalar types, this converts the value to this type.
    For composite types, this may construct an instance of this type.

    Args:
        val: The value to convert or use for construction

    Returns:
        The converted/constructed value
    """
    ...

fisinstance(other)

Check if other is an instance of this ftype.

Source code in src/finch/algebra/ftypes.py

def fisinstance(self, other) -> builtins.bool:
    """
    Check if `other` is an instance of this ftype.
    """
    return ftype(other) == self

FTyped

Bases: Generic[FT]

Abstract base class for objects that can be formatted.

Source code in src/finch/algebra/ftypes.py

class FTyped(Generic[FT]):
    """
    Abstract base class for objects that can be formatted.
    """

    @property
    @abstractmethod
    def ftype(self) -> FT:
        """
        The ftype of the object.
        """
        ...

ftype abstractmethod property

The ftype of the object.

FiberTensor dataclass

Bases: OverrideTensor

A class representing a tensor with fiber structure.

Attributes:

Name	Type	Description
lvl	Level	a fiber allocator that manages the fibers in the tensor.

Source code in src/finch/tensor/fiber_tensor.py

@dataclass
class FiberTensor(OverrideTensor):
    """
    A class representing a tensor with fiber structure.

    Attributes:
        lvl: a fiber allocator that manages the fibers in the tensor.
    """

    lvl: Level
    pos: np.integer = np.intp(0)
    _device: Any = None

    def __post_init__(self):
        self._device = normalize_device(self._device)

    def __repr__(self):
        return f"FiberTensor(lvl={self.lvl})"

    @property
    def ftype(self):
        """
        Returns the ftype of the fiber tensor, which is a FiberTensorFType.
        """
        return FiberTensorFType(self.lvl.ftype, self._device)

    def with_fill(self, fill_value: AbstractFill) -> FiberTensor:
        return replace(self, lvl=self.lvl.with_fill(fill_value))

    @property
    def shape(self):
        return self.lvl.shape

    @property
    def val(self):
        return self.lvl.val

    @property
    def ndim(self):
        return self.lvl.ndim

    @property
    def shape_type(self):
        return self.lvl.shape_type

    @property
    def element_type(self):
        return self.lvl.element_type

    @property
    def fill_value(self):
        return self.lvl.fill_value

    @property
    def device(self):
        return self._device

    def to_device(self, device, /, *, stream=None):
        if stream is not None:
            raise ValueError(f"stream argument is not supported; got {stream!r}")
        device = normalize_device(device)
        if device == self.device:
            return self
        return FiberTensor(self.lvl, self.pos, device)

    @property
    def position_type(self):
        return self.lvl.position_type

    @property
    def buffer_factory(self):
        return self.lvl.buffer_factory

    def to_numpy(self) -> np.ndarray:
        # TODO: temporary for dense only. TBD in sparse_level PR
        return np.reshape(self.lvl.val.arr, self.shape, copy=False)

    def to_scipy(self):
        from .level import (
            DenseLevel,
            ElementLevel,
            SparseCOOLevel,
            SparseListLevel,
        )

        if self.ndim != 2:
            raise ValueError("SciPy sparse arrays must be two-dimensional.")

        if self.fill_value != 0:
            raise ValueError("SciPy CSR conversion requires a zero fill value.")

        assert self.pos == 0

        match self.lvl:
            case SparseCOOLevel(
                lvl=ElementLevel() as element,
                tbl=(row, col),
            ):
                return sps.coo_array(
                    (element.val.arr, (row.arr, col.arr)),
                    shape=self.shape,
                    copy=False,
                )

            case DenseLevel(
                lvl=SparseListLevel(
                    lvl=ElementLevel() as element,
                    ptr=ptr,
                    idx=idx,
                )
            ):
                return sps.csr_array(
                    (element.val.arr, idx.arr, ptr.arr),
                    shape=self.shape,
                    copy=False,
                )

            case _:
                raise NotImplementedError(
                    f"Finch format {self.ftype} is not supported by SciPy conversion."
                )

    def item(self):
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to Python scalar.")
        return self.to_numpy().item()

    @classmethod
    def from_scipy_csr(
        cls,
        obj: Any,
        /,
        *,
        dtype=None,
        device=None,
        copy=None,
    ):
        from .level import (
            DenseLevel,
            ElementLevel,
            SparseListLevel,
            dense,
            element,
            sparse_list,
        )

        if dtype is not None:
            if copy is False:
                obj = obj.astype(dtype, copy=False)
            else:
                obj = obj.astype(dtype, copy=True)
        elif copy is True:
            obj = obj.copy()

        if not obj.has_canonical_format:
            if copy is False:
                raise ValueError(
                    "Unable to avoid copy while creating an array as requested."
                )
            obj = obj.copy()
            obj.sum_duplicates()

        data = obj.data
        indices = obj.indices
        indptr = obj.indptr

        element_lvl = element(
            fill_type=np.zeros((), dtype=data.dtype)[()],
            element_type=ftype(data.dtype),
            position_type=ftype(indices.dtype),
        )
        sparse_lvl = sparse_list(element_lvl, dimension_type=ftype(indices.dtype))
        dense_lvl = dense(sparse_lvl, dimension_type=ftype(indices.dtype))

        return cls(
            DenseLevel(
                SparseListLevel(
                    ElementLevel(
                        _format=element_lvl,
                        _val=NumpyBuffer(data),
                    ),
                    dimension=sparse_lvl.dimension_type(obj.shape[1]),
                    ptr=NumpyBuffer(indptr),
                    idx=NumpyBuffer(indices),
                ),
                dimension=dense_lvl.dimension_type(obj.shape[0]),
            ),
            _device=device,
        )

    @classmethod
    def from_scipy_coo(
        cls,
        obj: Any,
        /,
        *,
        dtype=None,
        device=None,
        copy=None,
    ):
        from .level import (
            ElementLevel,
            SparseCOOLevel,
            element,
        )

        if dtype is not None:
            if copy is False:
                obj = obj.astype(dtype, copy=False)
            else:
                obj = obj.astype(dtype, copy=True)
        elif copy is True:
            obj = obj.copy()

        if not obj.has_canonical_format:
            if copy is False:
                raise ValueError(
                    "Unable to avoid copy while creating an array as requested."
                )
            obj = obj.copy()
            obj.sum_duplicates()

        data = obj.data
        row = obj.row
        col = obj.col
        index_type = ftype(row.dtype)

        element_lvl = element(
            fill_type=np.zeros((), dtype=data.dtype)[()],
            element_type=ftype(data.dtype),
            position_type=index_type,
        )

        return cls(
            SparseCOOLevel(
                ElementLevel(
                    _format=element_lvl,
                    _val=NumpyBuffer(data),
                ),
                shape=tuple(index_type(d) for d in obj.shape),
                tbl=(NumpyBuffer(row), NumpyBuffer(col)),
            ),
            _device=device,
        )

ftype property

Returns the ftype of the fiber tensor, which is a FiberTensorFType.

FiberTensorFType dataclass

Bases: FinchTensorFType, ImmutableStructFType

An abstract base class representing the ftype of a fiber tensor.

Attributes:

Name	Type	Description
lvl_t	LevelFType	The level ftype to be used for the tensor.

Source code in src/finch/tensor/fiber_tensor.py

@dataclass(unsafe_hash=True)
class FiberTensorFType(FinchTensorFType, ImmutableStructFType):
    """
    An abstract base class representing the ftype of a fiber tensor.

    Attributes:
        lvl_t: The level ftype to be used for the tensor.
    """

    lvl_t: LevelFType
    _device: Any = None

    def __post_init__(self):
        self._device = normalize_device(self._device)

    @property
    def struct_name(self):
        # TODO: include dt = np.dtype(self.buf_t.element_type)
        return "FiberTensorFType"

    @property
    def struct_fields(self):
        return [
            ("lvl", self.lvl_t),
            ("shape", TupleFType.from_tuple(self.shape_type)),
            ("pos", self.position_type),
        ]

    def get_child_type(self, attr):
        if attr != "lvl":
            raise TypeError(f"{self} does not support child {attr!r}")
        return self.lvl_t

    def get_child(self, obj, attr):
        self.get_child_type(attr)
        return asm.GetAttr(obj, asm.Literal(attr))

    def with_fill(self, fill_type: AbstractFillFType) -> FiberTensorFType:
        """Rebuild this ftype with the leaf fill replaced."""
        return FiberTensorFType(self.lvl_t.with_fill(fill_type), self._device)

    def fisinstance(self, other):
        other_t = ftype(other)
        if is_dynamic(self.fill_type) and isinstance(other_t, FiberTensorFType):
            # A dynamic-fill ftype accepts any fill value of matching dtype.
            other_t = other_t.with_fill(self.fill_type)
        return other_t == self

    def construct(self, shape: tuple[int, ...], fill_value: Any = None) -> FiberTensor:
        """
        Creates an instance of a FiberTensor with the given arguments.
        """
        # The fill value may be raw, which is a static fill.
        fill = resolve_fill(
            self.fill_type, None if fill_value is None else as_fill(fill_value)
        )
        # The levels are built with a static fill, which the tensor then
        # replaces with its own.
        fmt = self.with_fill(StaticFillFType(fill.value))
        tns = FiberTensor(
            fmt.lvl_t.construct(shape=shape, pos=1),
            fmt.position_type(0),
            _device=fmt.device,
        )
        return tns.with_fill(fill)

    def __call__(self, val: Any) -> FiberTensor:
        """
        Convert a tensor to this fiber tensor type.

        Args:
            val: A tensor to convert to this type.
        Returns:
            A FiberTensor instance of this type.
        """
        raise NotImplementedError(
            f"Tensor conversion not yet implemented for {type(self).__name__}"
        )

    def __str__(self):
        return f"FiberTensorFType({self.lvl_t})"

    @property
    def ndim(self):
        return self.lvl_t.ndim

    @property
    def shape_type(self):
        return self.lvl_t.shape_type

    @property
    def element_type(self):
        return self.lvl_t.element_type

    @property
    def fill_type(self):
        return self.lvl_t.fill_type

    @property
    def device(self):
        return self._device

    @property
    def position_type(self):
        return self.lvl_t.position_type

    @property
    def buffer_factory(self):
        return self.lvl_t.buffer_factory

    @property
    def buffer_type(self):
        return self.lvl_t.buffer_type

    @property
    def level_format_properties(self) -> list[FormatProperty]:
        return self.lvl_t.level_format_properties(0)

    def unfurl(self, ctx, tns, ext, mode, proto):
        tns = ctx.resolve(tns)
        return self.lvl_t.level_unfurl(ctx, tns, ext, mode, proto, ctx(tns.pos))

    def lower_freeze(self, ctx, tns, op):
        return self.lvl_t.level_lower_freeze(
            ctx, ctx(tns.lvl), op, asm.Literal(self.position_type(1))
        )

    def lower_thaw(self, ctx, tns, op):
        return self.lvl_t.level_lower_thaw(
            ctx, ctx(tns.lvl), op, asm.Literal(self.position_type(1))
        )

    def lower_unwrap(self, ctx, obj):
        return self.lvl_t.level_lower_unwrap(ctx, obj, ctx(obj.pos))

    def lower_increment(self, ctx, obj, op, val):
        return self.lvl_t.level_lower_increment(ctx, obj, op, val, ctx(obj.pos))

    def lower_declare(self, ctx, tns, init, op, shape):
        return self.lvl_t.level_lower_declare(
            ctx, ctx(tns.lvl), init, op, shape, asm.Literal(self.position_type(1))
        )

    def lower_dim(self, ctx, obj, r):
        return self.lvl_t.level_lower_dim(ctx, ctx(obj.lvl), r)

    def lower_fill(self, ctx, obj):
        if is_static(self.fill_type):
            return asm.Literal(self.fill_type.value)
        return self.lvl_t.lower_fill(ctx(obj.lvl))

    def from_fields(self, *args) -> FiberTensor:
        lvl, shape, pos = args
        return FiberTensor(lvl, pos, self.device)

    def from_numpy(self, arr: np.ndarray, fill_value: Any = None) -> FiberTensor:
        # The fill value may be raw, which is a static fill.
        fill = resolve_fill(
            self.fill_type, None if fill_value is None else as_fill(fill_value)
        )

        def build_level(lvl_t, shape):
            match lvl_t:
                case levels.DenseLevelFType():
                    if not shape:
                        raise ValueError("Array rank does not match the fiber format")
                    return levels.DenseLevel(
                        build_level(lvl_t.lvl_t, shape[1:]),
                        lvl_t.dimension_type(shape[0]),
                    )
                case levels.ElementLevelFType():
                    if shape:
                        raise ValueError("Array rank does not match the fiber format")
                    return lvl_t.with_fill(fill.ftype).from_fields(arr, fill.value)
                case _:
                    raise NotImplementedError(
                        f"NumPy conversion does not support {type(lvl_t).__name__}"
                    )

        return FiberTensor(
            build_level(self.lvl_t, arr.shape),
            pos=self.position_type(0),
            _device=self.device,
        )

__call__(val)

Convert a tensor to this fiber tensor type.

Parameters:

Name	Type	Description	Default
val	Any	A tensor to convert to this type.	required

Returns: A FiberTensor instance of this type.

Source code in src/finch/tensor/fiber_tensor.py

def __call__(self, val: Any) -> FiberTensor:
    """
    Convert a tensor to this fiber tensor type.

    Args:
        val: A tensor to convert to this type.
    Returns:
        A FiberTensor instance of this type.
    """
    raise NotImplementedError(
        f"Tensor conversion not yet implemented for {type(self).__name__}"
    )

construct(shape, fill_value=None)

Creates an instance of a FiberTensor with the given arguments.

Source code in src/finch/tensor/fiber_tensor.py

def construct(self, shape: tuple[int, ...], fill_value: Any = None) -> FiberTensor:
    """
    Creates an instance of a FiberTensor with the given arguments.
    """
    # The fill value may be raw, which is a static fill.
    fill = resolve_fill(
        self.fill_type, None if fill_value is None else as_fill(fill_value)
    )
    # The levels are built with a static fill, which the tensor then
    # replaces with its own.
    fmt = self.with_fill(StaticFillFType(fill.value))
    tns = FiberTensor(
        fmt.lvl_t.construct(shape=shape, pos=1),
        fmt.position_type(0),
        _device=fmt.device,
    )
    return tns.with_fill(fill)

with_fill(fill_type)

Rebuild this ftype with the leaf fill replaced.

Source code in src/finch/tensor/fiber_tensor.py

def with_fill(self, fill_type: AbstractFillFType) -> FiberTensorFType:
    """Rebuild this ftype with the leaf fill replaced."""
    return FiberTensorFType(self.lvl_t.with_fill(fill_type), self._device)

FillTensor

Bases: OverrideTensor

A tensor that has a specific shape but contains no actual data. Used primarily for broadcasting operations where a tensor of a specific shape is needed but the values are irrelevant.

Source code in src/finch/tensor/patterns.py

class FillTensor(OverrideTensor):
    """
    A tensor that has a specific shape but contains no actual data.
    Used primarily for broadcasting operations where a tensor of a specific
    shape is needed but the values are irrelevant.
    """

    def with_fill(self, fill_value: AbstractFill) -> FillTensor:
        return FillTensor(self._shape, fill_value)

    def __init__(self, shape, fill_value):
        self._shape = shape
        self._fill = as_fill(fill_value)

    def __getitem__(self, idxs):
        return Scalar(self._fill.value, fill_value=self._fill)

    def item(self):
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to Python scalar.")
        return self._fill.value

    def to_numpy(self):
        raise NotImplementedError(f"{type(self).__name__} does not support to_numpy.")

    def to_scipy(self):
        raise NotImplementedError(f"{type(self).__name__} does not support to_scipy.")

    @property
    def shape(self):
        return self._shape

    @property
    def fill_value(self) -> Any:
        """Default fill value."""
        return self._fill.value

    @property
    def element_type(self) -> FType:
        """Data type of the tensor's elements."""
        return self.ftype.element_type

    @property
    def shape_type(self) -> tuple[FType, ...]:
        """Shape type of the tensor."""
        return self.ftype.shape_type

    @property
    def ftype(self):
        return FillTensorFType(
            self._fill.ftype,
            self._fill.element_type,
            tuple(ftype(dim) for dim in self.shape),
        )

element_type property

Data type of the tensor's elements.

fill_value property

Default fill value.

shape_type property

Shape type of the tensor.

IndexTensor

Bases: OverrideTensor

A tensor that has a specific shape and returns the linear (flattened) index used to access it.

Source code in src/finch/tensor/patterns.py

class IndexTensor(OverrideTensor):
    """
    A tensor that has a specific shape and returns the linear (flattened) index used
    to access it.
    """

    def __init__(
        self,
        shape,
        element_type: FType | type = np.intp,
        fill_value: AbstractFill | None = None,
    ):
        self._shape = tuple(shape)
        self._element_type = ftype(element_type)
        self._fill = fill_value

    def with_fill(self, fill_value: AbstractFill) -> IndexTensor:
        return IndexTensor(self._shape, self._element_type, fill_value)

    def __getitem__(self, idxs):
        if self.ndim == 0 and idxs in ((), Ellipsis, (...,)):
            return Scalar(self.fill_value, fill_value=self.fill_value)
        if not isinstance(idxs, tuple):
            idxs = (idxs,)
        idxs = tuple(idx for idx in idxs if idx is not Ellipsis)
        if len(idxs) != self.ndim:
            raise IndexError("Incorrect number of indices for IndexTensor.")
        flat_index = 0
        for idx, dim in zip(idxs, self.shape, strict=True):
            flat_index = flat_index * dim + idx
        flat_index = self._element_type(flat_index)
        return Scalar(flat_index, fill_value=self.fill_value)

    def item(self):
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to Python scalar.")
        return self.fill_value

    def to_numpy(self):
        raise NotImplementedError(f"{type(self).__name__} does not support to_numpy.")

    def to_scipy(self):
        raise NotImplementedError(f"{type(self).__name__} does not support to_scipy.")

    @property
    def shape(self):
        return self._shape

    @property
    def fill_value(self) -> Any:
        """Default fill value."""
        return resolve_fill(self.ftype.fill_type, self._fill).value

    @property
    def element_type(self) -> FType:
        """Data type of the tensor's elements."""
        return self.ftype.element_type

    @property
    def shape_type(self) -> tuple[FType, ...]:
        """Shape type of the tensor."""
        return self.ftype.shape_type

    @property
    def ftype(self):
        return IndexTensorFType(
            self._element_type,
            tuple(ftype(dim) for dim in self.shape),
            None if self._fill is None else self._fill.ftype,
        )

element_type property

Data type of the tensor's elements.

fill_value property

Default fill value.

shape_type property

Shape type of the tensor.

LazyTensor

Bases: OverrideTensor

Source code in src/finch/interface/lazy.py

class LazyTensor(OverrideTensor):
    def __init__(
        self,
        data: Alias,
        ctx: EffectBlob,
        shape: tuple,
        fill_value: AbstractFill,
        element_type: FType,
        device=None,
    ):
        self.data: Alias = data
        self.ctx = ctx
        self._shape = shape
        # Held as an AbstractFill so a dynamic fill survives ftype round trips.
        self._fill_value = fill_value
        self._element_type = element_type
        self._device = normalize_device(device)

    def override_module(self):
        return sys.modules[__name__]

    @property
    def ftype(self):
        return LazyTensorFType(
            _fill_type=self._fill_value.ftype,
            _element_type=self._element_type,
            _shape_type=tuple(ftype(dim) for dim in self._shape),
            _device=self._device,
        )

    @property
    def shape(self) -> tuple:
        """
        Returns the shape of the LazyTensor as a tuple.
        The shape is determined by the data and is a static property.
        """
        return self._shape

    @property
    def fill_value(self) -> Any:
        """Default value to fill the tensor."""
        return self._fill_value.value

    @property
    def device(self):
        return self._device

    def to_device(self, device, /, *, stream=None):
        if stream is not None:
            raise ValueError(f"stream argument is not supported; got {stream!r}")
        device = normalize_device(device)
        if device == self.device:
            return self
        return LazyTensor(
            self.data,
            self.ctx,
            self._shape,
            self._fill_value,
            self._element_type,
            device=device,
        )

    @property
    def element_type(self) -> FType:
        """Data type of the tensor elements."""
        return self.ftype.element_type

    @property
    def shape_type(self) -> tuple:
        """Shape type of the tensor."""
        return self.ftype.shape_type

    def item(self):
        raise ValueError(
            "Cannot convert LazyTensor to Python scalar. "
            "Use compute() to evaluate it first."
        )

    def to_numpy(self):
        raise ValueError(
            "Cannot convert LazyTensor to Python scalar. "
            "Use compute() to evaluate it first."
        )

    def to_scipy(self):
        raise ValueError(
            "Cannot convert LazyTensor to Python scalar. "
            "Use compute() to evaluate it first."
        )

    # raise ValueError for unsupported operations according to the data-apis spec.
    # NOT tested, since this isn't necessary as it will throw an error anyways.
    def __complex__(self) -> complex:
        """
        Converts the LazyTensor to a complex number.
        This is a no-op for LazyTensors, as they are symbolic representations.
        """
        raise ValueError(
            "Cannot convert LazyTensor to complex. Use compute() to evaluate it first."
        )

    def __float__(self) -> float:
        """
        Converts the LazyTensor to a float.
        This is a no-op for LazyTensors, as they are symbolic representations.
        """
        raise ValueError(
            "Cannot convert LazyTensor to float. Use compute() to evaluate it first."
        )

    def __int__(self) -> int:
        """
        Converts the LazyTensor to an int.
        This is a no-op for LazyTensors, as they are symbolic representations.
        """
        raise ValueError(
            "Cannot convert LazyTensor to int. Use compute() to evaluate it first."
        )

    def __bool__(self) -> bool:
        """
        Converts the LazyTensor to a bool.
        This is a no-op for LazyTensors, as they are symbolic representations.
        """
        raise ValueError(
            "Cannot convert LazyTensor to bool. Use compute() to evaluate it first."
        )

    def reshape(self, shape, /, *, copy=None):
        return reshape(self, shape, copy=copy)

element_type property

Data type of the tensor elements.

fill_value property

Default value to fill the tensor.

shape property

Returns the shape of the LazyTensor as a tuple. The shape is determined by the data and is a static property.

shape_type property

Shape type of the tensor.

__bool__()

Converts the LazyTensor to a bool. This is a no-op for LazyTensors, as they are symbolic representations.

Source code in src/finch/interface/lazy.py

def __bool__(self) -> bool:
    """
    Converts the LazyTensor to a bool.
    This is a no-op for LazyTensors, as they are symbolic representations.
    """
    raise ValueError(
        "Cannot convert LazyTensor to bool. Use compute() to evaluate it first."
    )

__complex__()

Converts the LazyTensor to a complex number. This is a no-op for LazyTensors, as they are symbolic representations.

Source code in src/finch/interface/lazy.py

def __complex__(self) -> complex:
    """
    Converts the LazyTensor to a complex number.
    This is a no-op for LazyTensors, as they are symbolic representations.
    """
    raise ValueError(
        "Cannot convert LazyTensor to complex. Use compute() to evaluate it first."
    )

__float__()

Converts the LazyTensor to a float. This is a no-op for LazyTensors, as they are symbolic representations.

Source code in src/finch/interface/lazy.py

def __float__(self) -> float:
    """
    Converts the LazyTensor to a float.
    This is a no-op for LazyTensors, as they are symbolic representations.
    """
    raise ValueError(
        "Cannot convert LazyTensor to float. Use compute() to evaluate it first."
    )

__int__()

Converts the LazyTensor to an int. This is a no-op for LazyTensors, as they are symbolic representations.

Source code in src/finch/interface/lazy.py

def __int__(self) -> int:
    """
    Converts the LazyTensor to an int.
    This is a no-op for LazyTensors, as they are symbolic representations.
    """
    raise ValueError(
        "Cannot convert LazyTensor to int. Use compute() to evaluate it first."
    )

NumPyFType

Bases: TensorFType

Source code in src/finch/tensor/np_wrapper.py

class NumPyFType(TensorFType):
    def __init__(self, dtype: np.dtype, ndim: int, device=None):
        self._dtype = np.dtype(dtype)
        self._ndim = ndim
        self._device = normalize_device(device)

    @property
    def fill_type(self) -> AbstractFillFType:
        return StaticFillFType(self._dtype.type(0))

    @property
    def element_type(self) -> FType:
        return ftype(self._dtype.type)

    @property
    def shape_type(self) -> tuple[FType, ...]:
        return (intp,) * self._ndim

    @property
    def device(self):
        return self._device

    def construct(self, shape: tuple) -> "NumPyWrapper":
        # creates a zero-filled tensor
        return NumPyWrapper(np.zeros(shape, dtype=self._dtype), device=self.device)

    def __call__(self, val: "NumPyWrapper") -> "NumPyWrapper":
        """
        Convert a tensor to this numpy tensor type.

        Args:
            val: A tensor to convert to this type.
        Returns:
            A NumPyWrapper instance of this type.
        """
        raise NotImplementedError(
            f"Tensor conversion not yet implemented for {type(self).__name__}"
        )

    def __eq__(self, other):
        if not isinstance(other, NumPyFType):
            return False
        return (
            self._dtype == other._dtype
            and self._ndim == other._ndim
            and self.device == other.device
        )

    def __hash__(self):
        return hash((self._dtype, self._ndim, self.device))

    def from_numpy(self, arr: np.ndarray) -> "NumPyWrapper":
        return NumPyWrapper(arr, device=self.device)

__call__(val)

Convert a tensor to this numpy tensor type.

Parameters:

Name	Type	Description	Default
val	NumPyWrapper	A tensor to convert to this type.	required

Returns: A NumPyWrapper instance of this type.

Source code in src/finch/tensor/np_wrapper.py

def __call__(self, val: "NumPyWrapper") -> "NumPyWrapper":
    """
    Convert a tensor to this numpy tensor type.

    Args:
        val: A tensor to convert to this type.
    Returns:
        A NumPyWrapper instance of this type.
    """
    raise NotImplementedError(
        f"Tensor conversion not yet implemented for {type(self).__name__}"
    )

NumPyWrapper

Bases: Tensor

Source code in src/finch/tensor/np_wrapper.py

class NumPyWrapper(Tensor):
    def __init__(self, data: np.ndarray, device=None):
        self._data = np.asarray(data)
        self._device = normalize_device(device)

    @property
    def ftype(self) -> NumPyFType:
        return NumPyFType(self._data.dtype, self._data.ndim, self._device)

    @property
    def shape(self) -> tuple:
        return self._data.shape

    @property
    def fill_value(self) -> Any:
        """Default fill value."""
        return resolve_fill(self.ftype.fill_type).value

    @property
    def element_type(self) -> FType:
        """Data type of the tensor's elements."""
        return self.ftype.element_type

    @property
    def shape_type(self) -> tuple[FType, ...]:
        """Shape type of the tensor."""
        return self.ftype.shape_type

    @property
    def device(self):
        return self._device

    def to_device(self, device, /, *, stream=None):
        if stream is not None:
            raise ValueError(f"stream argument is not supported; got {stream!r}")
        device = normalize_device(device)
        if device == self.device:
            return self
        return NumPyWrapper(self._data, device=device)

    def item(self):
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to Python scalar.")
        return self._data.item()

    def to_numpy(self):
        return self._data

    def to_scipy(self):
        raise NotImplementedError(f"{type(self).__name__} does not support to_scipy.")

element_type property

Data type of the tensor's elements.

fill_value property

Default fill value.

shape_type property

Shape type of the tensor.

NumpyBuffer

Bases: Buffer

A buffer that uses NumPy arrays to store data. This is a concrete implementation of the Buffer class.

Source code in src/finch/codegen/buffers/numpy_buffer.py

class NumpyBuffer(Buffer):
    """
    A buffer that uses NumPy arrays to store data. This is a concrete implementation
    of the Buffer class.
    """

    def __init__(self, arr: np.ndarray):
        if not arr.flags["C_CONTIGUOUS"]:
            raise ValueError("NumPy array must be C-contiguous")
        self.arr = arr

    @property
    def ftype(self):
        """
        Returns the ftype of the buffer, which is a NumpyBufferFType.
        """
        return NumpyBufferFType(ftype(self.arr.dtype))

    # TODO should be property
    def length(self):
        return self.arr.size

    def load(self, index: int):
        value = self.arr[index]
        if isinstance(self.ftype.element_type, TupleFType):
            return tuple(
                value[name] for name in self.ftype.element_type.struct_fieldnames
            )
        return value

    def store(self, index: int, value):
        self.arr[index] = value

    def resize(self, new_length: int):
        self.arr = np.resize(self.arr, new_length)

    def __str__(self):
        arr_str = str(self.arr).replace("\n", "")
        return f"np_buf({arr_str})"

    def __repr__(self):
        arr_repr = repr(self.arr).replace("\n", "")
        return f"NumpyBuffer({arr_repr})"

ftype property

Returns the ftype of the buffer, which is a NumpyBufferFType.

NumpyBufferFType

Bases: CBufferFType, NumbaBufferFType, MLIRBufferFType, CUnpackableFType

A ftype for buffers that uses NumPy arrays. This is a concrete implementation of the BufferFType class.

Source code in src/finch/codegen/buffers/numpy_buffer.py

class NumpyBufferFType(
    CBufferFType, NumbaBufferFType, MLIRBufferFType, CUnpackableFType
):
    """
    A ftype for buffers that uses NumPy arrays. This is a concrete implementation
    of the BufferFType class.
    """

    def __init__(self, element_type: FType):
        self._element_type = ftype(np_dtype(ftype(element_type)))

    @property
    def _dtype(self):
        return np_dtype(self._element_type)

    def __eq__(self, other):
        if not isinstance(other, NumpyBufferFType):
            return False
        return self._element_type == other._element_type

    def __str__(self):
        return f"np_buf_t({qual_str(self._dtype.type)})"

    def __repr__(self):
        return f"NumpyBufferFType({repr(self._element_type)})"

    @property
    def length_type(self):
        """
        Returns the type used for the length of the buffer.
        """
        return ftypes.intp

    @property
    def element_type(self):
        """
        Returns the type of elements stored in the buffer.
        This is typically the same as the dtype used to create the buffer.
        """
        return self._element_type

    def __hash__(self):
        return hash(self._element_type)

    def __call__(self, len: int = 0, element_type: FType | None = None):
        return NumpyBuffer(np.zeros(len, dtype=self._dtype))

    def c_type(self):
        return ctypes.POINTER(CNumpyBuffer)

    def c_length(self, ctx: "CContext", buf: CBufferFields):
        return buf.length

    def c_data(self, ctx: "CContext", buf: CBufferFields):
        return buf.data

    def c_load(self, ctx: "CContext", buf: CBufferFields, idx: "AssemblyExpression"):
        return f"({buf.data})[{ctx(idx)}]"

    def c_store(
        self,
        ctx: "CContext",
        buf: CBufferFields,
        idx: "AssemblyExpression",
        value: "AssemblyExpression",
    ):
        ctx.exec(f"{ctx.feed}({buf.data})[{ctx(idx)}] = {ctx(value)};")

    def c_resize(self, ctx, buf: CBufferFields, new_len):
        new_len = ctx(ctx.cache("len", new_len))
        data = buf.data
        length = buf.length
        obj = buf.obj
        t = ctx.ctype_name(c_type(self.element_type))
        ctx.exec(
            f"{ctx.feed}{data} = ({t}*){obj}->resize(&{obj}->arr, {new_len});\n"
            f"{ctx.feed}{length} = {new_len};"
        )
        return

    def c_unpack(self, ctx, var_n, val):
        """
        Unpack the buffer into C context.
        """
        data = ctx.freshen(var_n, "data")
        length = ctx.freshen(var_n, "length")
        t = ctx.ctype_name(c_type(self.element_type))
        ctx.add_header("#include <stddef.h>")
        ctx.exec(
            f"{ctx.feed}{t}* {data} = ({t}*){ctx(val)}->data;\n"
            f"{ctx.feed}size_t {length} = {ctx(val)}->length;"
        )

        return CBufferFields(data, length, var_n)

    def c_repack(self, ctx, lhs, obj):
        """
        Repack the buffer from C context.
        """
        ctx.exec(
            f"{ctx.feed}{lhs}->data = (void*){obj.data};\n"
            f"{ctx.feed}{lhs}->length = {obj.length};"
        )
        return

    def serialize_to_c(self, obj):
        """
        Serialize the NumPy buffer to a C-compatible structure.
        """
        data = ctypes.c_void_p(obj.arr.ctypes.data)
        length = obj.arr.size
        obj._self_obj = ctypes.py_object(obj)
        obj._c_callback = numpy_buffer_resize_callback
        obj._c_buffer = CNumpyBuffer(obj._self_obj, data, length, obj._c_callback)
        return ctypes.pointer(obj._c_buffer)

    def deserialize_from_c(self, obj, c_buffer):
        """
        Update this buffer based on how the C call modified the CNumpyBuffer structure.
        """
        # this is handled by the resize callback

    def construct_from_c(self, c_buffer):
        """
        Construct a NumpyBuffer from a C-compatible structure.
        """
        return c_buffer.contents.arr

    def numba_type(self) -> type:
        return list[np.ndarray]

    def numba_jitclass_type(self) -> numba.types.Type:
        return numba.types.ListType(
            numba.types.Array(numba.from_dtype(self._dtype), 1, "C")
        )

    def numba_length(self, ctx, buf: NumbaBufferFields):
        arr = buf.arr
        return f"len({arr})"

    def numba_load(self, ctx, buf, idx):
        buf = cast(NumbaBufferFields, buf)
        arr = buf.arr
        if isinstance(self.element_type, TupleFType):
            idx = ctx(ctx.cache("idx", idx))
            fields = ", ".join(
                f"{arr}[{idx}]['{name}']"
                for name in self.element_type.struct_fieldnames
            )
            return f"({fields},)"
        return f"{arr}[{ctx(idx)}]"

    def numba_store(self, ctx, buf, idx, value=None):
        buf = cast(NumbaBufferFields, buf)
        arr = buf.arr
        if isinstance(self.element_type, TupleFType):
            idx = ctx(ctx.cache("idx", idx))
            value = ctx.cache("val", value)
            val_code = ctx(value)
            for i, name in enumerate(self.element_type.struct_fieldnames):
                ctx.exec(f"{ctx.feed}{arr}[{idx}]['{name}'] = {val_code}[{i}]")
            return
        ctx.exec(f"{ctx.feed}{arr}[{ctx(idx)}] = {ctx(value)}")

    def numba_resize(self, ctx, buf: NumbaBufferFields, new_len):
        arr = buf.arr
        ctx.exec(f"{ctx.feed}{arr} = numpy.resize({arr}, {ctx(new_len)})")

    def numba_unpack(self, ctx, var_n, val):
        """
        Unpack the buffer into Numba context.
        """
        arr = ctx.freshen(var_n, "arr")
        ctx.exec(f"{ctx.feed}{arr} = {ctx(val)}[0]")

        return NumbaBufferFields(arr, var_n)

    def numba_repack(self, ctx, lhs, obj):
        """
        Repack the buffer from Numba context.
        """
        ctx.exec(f"{ctx.feed}{lhs}[0] = {obj.arr}")
        return

    def serialize_to_numba(self, obj):
        """
        Serialize the NumPy buffer to a Numba-compatible object.
        """
        return numba.typed.List([obj.arr])

    def deserialize_from_numba(self, obj, numba_buffer):
        obj.arr = numba_buffer[0]
        return

    def construct_from_numba(self, numba_buffer):
        """
        Construct a NumpyBuffer from a Numba-compatible object.
        """
        return NumpyBuffer(numba_buffer[0])

    def mlir_type(self):
        return "!llvm.ptr"

    def mlir_buffer_type(self):
        return f"memref<?x{mlir_type(self.element_type)}>"

    # this is the rank-1 memref descriptor type
    def mlir_descriptor_type(self):
        return "!llvm.struct<(ptr, ptr, i64, array<1 x i64>, array<1 x i64>)>"

    # this is the callback buffer type
    def llvm_callback_type(self):
        return "!llvm.struct<(ptr, ptr, i64, ptr)>"

    # Create our own rank-1 memeref descriptor from a data pointer and length.
    def mlir_memref_from_pointer(self, ctx, data, length):
        memref_t = self.mlir_buffer_type()
        desc_t = self.mlir_descriptor_type()

        desc = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{desc} = llvm.mlir.undef : {desc_t}")

        for id, val in (
            ("0", data),
            ("1", data),
            ("2", ctx.constant(0, "i64")),
            ("3, 0", length),
            ("4, 0", ctx.constant(1, "i64")),
        ):
            fresh = ctx.new_ssa()
            ctx.exec(
                f"{ctx.feed}{fresh} = llvm.insertvalue {val}, {desc}[{id}] : {desc_t}"
            )
            desc = fresh

        buf = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{buf} = builtin.unrealized_conversion_cast "
            f"{desc} : {desc_t} to {memref_t}"
        )
        return buf

    # load the length of the NumPy buffer
    def mlir_length(self, ctx: "MLIRContext", buf: MLIRBufferFields):
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{desc} = llvm.load {buf.box} : !llvm.ptr -> {desc_t}")
        buffer = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{buffer} = builtin.unrealized_conversion_cast "
            f"{desc} : {desc_t} to {self.mlir_buffer_type()}"
        )
        c0 = ctx.constant(0, "index")
        res = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{res} = memref.dim {buffer}, {c0} : {self.mlir_buffer_type()}"
        )
        return res

    # Convert index values to mlir index types
    def mlir_index(self, ctx: "MLIRContext", idx: "AssemblyExpression"):
        i = ctx(idx)
        t = mlir_type(idx.result_type)
        if t == "index":
            return i
        cast_ = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{cast_} = arith.index_cast {i} : {t} to index")
        return cast_

    # Load buffer values of NumPy buffer
    def mlir_load(
        self, ctx: "MLIRContext", buf: MLIRBufferFields, idx: "AssemblyExpression"
    ):
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{desc} = llvm.load {buf.box} : !llvm.ptr -> {desc_t}")
        buffer = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{buffer} = builtin.unrealized_conversion_cast "
            f"{desc} : {desc_t} to {self.mlir_buffer_type()}"
        )
        i = self.mlir_index(ctx, idx)
        c0 = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{c0} = memref.load {buffer}[{i}] : {self.mlir_buffer_type()}"
        )
        return c0

    # Store values in the buffer
    def mlir_store(
        self,
        ctx: "MLIRContext",
        buf: Any,
        idx: Any,
        value: Any = None,
    ):
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{desc} = llvm.load {buf.box} : !llvm.ptr -> {desc_t}")
        buffer = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{buffer} = builtin.unrealized_conversion_cast "
            f"{desc} : {desc_t} to {self.mlir_buffer_type()}"
        )
        new = ctx(value)
        i = self.mlir_index(ctx, idx)
        ctx.exec(
            f"{ctx.feed}memref.store {new}, {buffer}[{i}] : {self.mlir_buffer_type()}"
        )

    # Resize the NumPy buffer and update its memref descriptor.
    def mlir_resize(self, ctx, buf: MLIRBufferFields, new_len):
        len_idx = mlir_cast_value(ctx, ctx(new_len), new_len.result_type, ftypes.intp)
        len = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{len} = arith.index_cast {len_idx} : index to i64")
        arr_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{arr_ptr} = llvm.getelementptr {buf.obj}[0, 0] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )

        cb_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{cb_ptr} = llvm.getelementptr {buf.obj}[0, 3] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )
        cb = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{cb} = llvm.load {cb_ptr} : !llvm.ptr -> !llvm.ptr")

        data = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{data} = llvm.call {cb}({arr_ptr}, {len}) "
            f": !llvm.ptr, (!llvm.ptr, i64) -> !llvm.ptr"
        )

        buffer = self.mlir_memref_from_pointer(ctx, data, len)
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{desc} = builtin.unrealized_conversion_cast "
            f"{buffer} : {self.mlir_buffer_type()} to {desc_t}"
        )
        ctx.exec(f"{ctx.feed}llvm.store {desc}, {buf.box} : {desc_t}, !llvm.ptr")

    # Unpack the NumPy buffer and create its memeref descriptor
    def mlir_unpack(self, ctx: MLIRContext, _var_n, val):
        obj = ctx(val)

        data_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{data_ptr} = llvm.getelementptr {obj}[0, 1] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )
        data = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{data} = llvm.load {data_ptr} : !llvm.ptr -> !llvm.ptr")

        length_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{length_ptr} = llvm.getelementptr {obj}[0, 2] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )
        length = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{length} = llvm.load {length_ptr} : !llvm.ptr -> i64")

        buffer = self.mlir_memref_from_pointer(ctx, data, length)
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{desc} = builtin.unrealized_conversion_cast "
            f"{buffer} : {self.mlir_buffer_type()} to {desc_t}"
        )
        box = ctx.new_ssa()
        count = ctx.constant(1, "i64")
        ctx.exec(
            f"{ctx.feed}{box} = llvm.alloca {count} x {desc_t} : (i64) -> !llvm.ptr"
        )
        ctx.exec(f"{ctx.feed}llvm.store {desc}, {box} : {desc_t}, !llvm.ptr")

        return MLIRBufferFields(box, obj)

    # Update the buffer wrapper with the resized pointer and length.
    def mlir_repack(self, ctx: MLIRContext, _lhs, obj):
        memref_t = self.mlir_buffer_type()
        desc_t = self.mlir_descriptor_type()
        desc = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{desc} = llvm.load {obj.box} : !llvm.ptr -> {desc_t}")
        buffer = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{buffer} = builtin.unrealized_conversion_cast "
            f"{desc} : {desc_t} to {memref_t}"
        )

        len_ssa = self.mlir_length(ctx, obj)
        len_i64 = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{len_i64} = arith.index_cast {len_ssa} : index to i64")
        len_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{len_ptr} = llvm.getelementptr {obj.obj}[0, 2] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )
        ctx.exec(f"{ctx.feed}llvm.store {len_i64}, {len_ptr} : i64, !llvm.ptr")

        data_ssa = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{data_ssa} = "
            f"memref.extract_aligned_pointer_as_index {buffer} "
            f": {memref_t} -> index"
        )
        data_i64 = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{data_i64} = arith.index_cast {data_ssa} : index to i64")
        data_llvm = ctx.new_ssa()
        ctx.exec(f"{ctx.feed}{data_llvm} = llvm.inttoptr {data_i64} : i64 to !llvm.ptr")
        data_ptr = ctx.new_ssa()
        ctx.exec(
            f"{ctx.feed}{data_ptr} = llvm.getelementptr {obj.obj}[0, 1] "
            f": (!llvm.ptr) -> !llvm.ptr, {self.llvm_callback_type()}"
        )
        ctx.exec(f"{ctx.feed}llvm.store {data_llvm}, {data_ptr} : !llvm.ptr, !llvm.ptr")

    # Package the NumPy array, data pointer, length, and resize callback for MLIR
    def serialize_to_mlir(self, obj):
        result = MLIRNumpyBuffer(
            ctypes.py_object(obj),
            ctypes.c_void_p(obj.arr.ctypes.data),
            obj.arr.size,
            numpy_buffer_resize_callback,
        )
        return ctypes.cast(ctypes.pointer(result), ctypes.c_void_p)

    def deserialize_from_mlir(self, obj, mlir_buffer):
        # The resize callback updates the owning NumPy buffer directly.
        pass

    def construct_from_mlir(self, mlir_buffer):
        # zero copy
        result = ctypes.cast(mlir_buffer, ctypes.POINTER(MLIRNumpyBuffer))
        return result.contents.arr

element_type property

Returns the type of elements stored in the buffer. This is typically the same as the dtype used to create the buffer.

length_type property

Returns the type used for the length of the buffer.

c_repack(ctx, lhs, obj)

Repack the buffer from C context.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def c_repack(self, ctx, lhs, obj):
    """
    Repack the buffer from C context.
    """
    ctx.exec(
        f"{ctx.feed}{lhs}->data = (void*){obj.data};\n"
        f"{ctx.feed}{lhs}->length = {obj.length};"
    )
    return

c_unpack(ctx, var_n, val)

Unpack the buffer into C context.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def c_unpack(self, ctx, var_n, val):
    """
    Unpack the buffer into C context.
    """
    data = ctx.freshen(var_n, "data")
    length = ctx.freshen(var_n, "length")
    t = ctx.ctype_name(c_type(self.element_type))
    ctx.add_header("#include <stddef.h>")
    ctx.exec(
        f"{ctx.feed}{t}* {data} = ({t}*){ctx(val)}->data;\n"
        f"{ctx.feed}size_t {length} = {ctx(val)}->length;"
    )

    return CBufferFields(data, length, var_n)

construct_from_c(c_buffer)

Construct a NumpyBuffer from a C-compatible structure.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def construct_from_c(self, c_buffer):
    """
    Construct a NumpyBuffer from a C-compatible structure.
    """
    return c_buffer.contents.arr

construct_from_numba(numba_buffer)

Construct a NumpyBuffer from a Numba-compatible object.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def construct_from_numba(self, numba_buffer):
    """
    Construct a NumpyBuffer from a Numba-compatible object.
    """
    return NumpyBuffer(numba_buffer[0])

deserialize_from_c(obj, c_buffer)

Update this buffer based on how the C call modified the CNumpyBuffer structure.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def deserialize_from_c(self, obj, c_buffer):
    """
    Update this buffer based on how the C call modified the CNumpyBuffer structure.
    """

numba_repack(ctx, lhs, obj)

Repack the buffer from Numba context.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def numba_repack(self, ctx, lhs, obj):
    """
    Repack the buffer from Numba context.
    """
    ctx.exec(f"{ctx.feed}{lhs}[0] = {obj.arr}")
    return

numba_unpack(ctx, var_n, val)

Unpack the buffer into Numba context.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def numba_unpack(self, ctx, var_n, val):
    """
    Unpack the buffer into Numba context.
    """
    arr = ctx.freshen(var_n, "arr")
    ctx.exec(f"{ctx.feed}{arr} = {ctx(val)}[0]")

    return NumbaBufferFields(arr, var_n)

serialize_to_c(obj)

Serialize the NumPy buffer to a C-compatible structure.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def serialize_to_c(self, obj):
    """
    Serialize the NumPy buffer to a C-compatible structure.
    """
    data = ctypes.c_void_p(obj.arr.ctypes.data)
    length = obj.arr.size
    obj._self_obj = ctypes.py_object(obj)
    obj._c_callback = numpy_buffer_resize_callback
    obj._c_buffer = CNumpyBuffer(obj._self_obj, data, length, obj._c_callback)
    return ctypes.pointer(obj._c_buffer)

serialize_to_numba(obj)

Serialize the NumPy buffer to a Numba-compatible object.

Source code in src/finch/codegen/buffers/numpy_buffer.py

def serialize_to_numba(self, obj):
    """
    Serialize the NumPy buffer to a Numba-compatible object.
    """
    return numba.typed.List([obj.arr])

OverrideTensor

Bases: Tensor

Source code in src/finch/tensor/override_tensor.py

class OverrideTensor(Tensor):
    @property
    def mT(self):
        return self.mod.matrix_transpose(self)

    @property
    def T(self):
        if self.ndim != 2:
            raise ValueError("T is only defined for 2D arrays")
        return self.mod.matrix_transpose(self)

    @property
    def mod(self):
        return self.override_module()

    def override_module(self):
        return sys.modules["finch.interface"]

    def __array_function__(self, func, types, args, kwargs):
        """Override NumPy functions using the __array_function__ protocol."""
        # https://numpy.org/neps/nep-0018-array-function-protocol.html
        override_func = getattr(self.mod, func.__name__, None)
        if override_func is None:
            return NotImplemented
        return override_func(*args, **kwargs)

    def __array_ufunc__(self, ufunc: np.ufunc, method, *inputs, **kwargs):
        """Override NumPy ufuncs using the __array_ufunc__ protocol."""
        # https://numpy.org/devdocs/user/basics.ufuncs.html#ufuncs-basics
        # https://numpy.org/devdocs/reference/ufuncs.html#ufuncs-methods
        if kwargs.get("out") is not None:
            raise NotImplementedError("out parameter is not supported")
        if kwargs.get("where") is not None:
            raise NotImplementedError("where parameter is not supported")
        if kwargs.get("casting") is not None:
            raise NotImplementedError("casting parameter is not supported")
        if kwargs.get("order") is not None:
            raise NotImplementedError("order parameter is not supported")
        if kwargs.get("axes") is not None:
            kwargs["axis"] = kwargs.pop("axes")
        if ufunc in element_wise_ufunc_map:
            if method == "__call__":
                return self.mod.elementwise(
                    element_wise_ufunc_map[ufunc], *inputs, **kwargs
                )
            if method == "reduce":
                return self.mod.reduce(ufunc, *inputs, **kwargs)
        if ufunc in ufunc_map:
            func_name = ufunc_map[ufunc]
            if method == "__call__":
                return getattr(self.mod, func_name)(*inputs, **kwargs)
        return NotImplemented

    def __array_namespace__(self, *, api_version=None):
        # https://data-apis.org/array-api/latest/API_specification/generated/array_api.array.__array_namespace__.html#array_api.array.__array_namespace__
        if api_version is None:
            api_version = "2024.12"

        if api_version not in {"2024.12"}:
            raise ValueError(f'"{api_version}" Array API version not supported.')

        return sys.modules["finch"]

    def __add__(self, other):
        return self.mod.add(self, other)

    def __radd__(self, other):
        return self.mod.add(other, self)

    def __sub__(self, other):
        return self.mod.subtract(self, other)

    def __rsub__(self, other):
        return self.mod.subtract(other, self)

    def __mul__(self, other):
        return self.mod.multiply(self, other)

    def __rmul__(self, other):
        return self.mod.multiply(other, self)

    def __abs__(self):
        return self.mod.abs(self)

    def __pos__(self):
        return self.mod.positive(self)

    def __neg__(self):
        return self.mod.negative(self)

    def __invert__(self):
        return self.mod.bitwise_invert(self)

    def __and__(self, other):
        return self.mod.bitwise_and(self, other)

    def __rand__(self, other):
        return self.mod.bitwise_and(other, self)

    def __lshift__(self, other):
        return self.mod.bitwise_left_shift(self, other)

    def __rlshift__(self, other):
        return self.mod.bitwise_left_shift(other, self)

    def __or__(self, other):
        return self.mod.bitwise_or(self, other)

    def __ror__(self, other):
        return self.mod.bitwise_or(other, self)

    def __rshift__(self, other):
        return self.mod.bitwise_right_shift(self, other)

    def __rrshift__(self, other):
        return self.mod.bitwise_right_shift(other, self)

    def __xor__(self, other):
        return self.mod.bitwise_xor(self, other)

    def __rxor__(self, other):
        return self.mod.bitwise_xor(other, self)

    def __truediv__(self, other):
        return self.mod.truediv(self, other)

    def __rtruediv__(self, other):
        return self.mod.truediv(other, self)

    def __floordiv__(self, other):
        return self.mod.floor_divide(self, other)

    def __rfloordiv__(self, other):
        return self.mod.floor_divide(other, self)

    def __mod__(self, other):
        return self.mod.mod(self, other)

    def __rmod__(self, other):
        return self.mod.mod(other, self)

    def __pow__(self, other):
        return self.mod.power(self, other)

    def __rpow__(self, other):
        return self.mod.power(other, self)

    def __matmul__(self, other):
        return self.mod.matmul(self, other)

    def __rmatmul__(self, other):
        return self.mod.matmul(other, self)

    def __sin__(self):
        return self.mod.sin(self)

    def __sinh__(self):
        return self.mod.sinh(self)

    def __cos__(self):
        return self.mod.cos(self)

    def __cosh__(self):
        return self.mod.cosh(self)

    def __tan__(self):
        return self.mod.tan(self)

    def __tanh__(self):
        return self.mod.tanh(self)

    def __asin__(self):
        return self.mod.asin(self)

    def __asinh__(self):
        return self.mod.asinh(self)

    def __acos__(self):
        return self.mod.acos(self)

    def __acosh__(self):
        return self.mod.acosh(self)

    def __atan__(self):
        return self.mod.atan(self)

    def __atanh__(self):
        return self.mod.atanh(self)

    def __atan2__(self, other):
        return self.mod.atan2(self, other)

    def __complex__(self):
        """
        Converts a zero-dimensional array to a Python `complex` object.
        """
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to complex.")
        # dispatch to the scalar value's `__complex__` method
        return complex(self.item())

    def __float__(self):
        """
        Converts a zero-dimensional array to a Python `float` object.
        """
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to float.")
        # dispatch to the scalar value's `__float__` method
        return float(self.item())

    def __int__(self):
        """
        Converts a zero-dimensional array to a Python `int` object.
        """
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to int.")
        # dispatch to the scalar value's `__int__` method
        return int(self.item())

    def __bool__(self):
        """
        Converts a zero-dimensional array to a Python `bool` object.
        """
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to bool.")
        # dispatch to the scalar value's `__bool__` method
        return bool(self.item())

    def __index__(self) -> int:
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to index.")
        return operator.index(self.item())

    def __log__(self):
        return self.mod.log(self)

    def __log1p__(self):
        return self.mod.log1p(self)

    def __log2__(self):
        return self.mod.log2(self)

    def __log10__(self):
        return self.mod.log10(self)

    def __logaddexp__(self, other):
        return self.mod.logaddexp(self, other)

    def __logical_and__(self, other):
        return self.mod.logical_and(self, other)

    def __logical_or__(self, other):
        return self.mod.logical_or(self, other)

    def __logical_xor__(self, other):
        return self.mod.logical_xor(self, other)

    def __logical_not__(self):
        return self.mod.logical_not(self)

    def __lt__(self, other):
        return self.mod.less(self, other)

    def __le__(self, other):
        return self.mod.less_equal(self, other)

    def __gt__(self, other):
        return self.mod.greater(self, other)

    def __ge__(self, other):
        return self.mod.greater_equal(self, other)

    def __eq__(self, other):
        return self.mod.equal(self, other)

    def __ne__(self, other):
        return self.mod.not_equal(self, other)

__array_function__(func, types, args, kwargs)

Override NumPy functions using the array_function protocol.

Source code in src/finch/tensor/override_tensor.py

def __array_function__(self, func, types, args, kwargs):
    """Override NumPy functions using the __array_function__ protocol."""
    # https://numpy.org/neps/nep-0018-array-function-protocol.html
    override_func = getattr(self.mod, func.__name__, None)
    if override_func is None:
        return NotImplemented
    return override_func(*args, **kwargs)

__array_ufunc__(ufunc, method, *inputs, **kwargs)

Override NumPy ufuncs using the array_ufunc protocol.

Source code in src/finch/tensor/override_tensor.py

def __array_ufunc__(self, ufunc: np.ufunc, method, *inputs, **kwargs):
    """Override NumPy ufuncs using the __array_ufunc__ protocol."""
    # https://numpy.org/devdocs/user/basics.ufuncs.html#ufuncs-basics
    # https://numpy.org/devdocs/reference/ufuncs.html#ufuncs-methods
    if kwargs.get("out") is not None:
        raise NotImplementedError("out parameter is not supported")
    if kwargs.get("where") is not None:
        raise NotImplementedError("where parameter is not supported")
    if kwargs.get("casting") is not None:
        raise NotImplementedError("casting parameter is not supported")
    if kwargs.get("order") is not None:
        raise NotImplementedError("order parameter is not supported")
    if kwargs.get("axes") is not None:
        kwargs["axis"] = kwargs.pop("axes")
    if ufunc in element_wise_ufunc_map:
        if method == "__call__":
            return self.mod.elementwise(
                element_wise_ufunc_map[ufunc], *inputs, **kwargs
            )
        if method == "reduce":
            return self.mod.reduce(ufunc, *inputs, **kwargs)
    if ufunc in ufunc_map:
        func_name = ufunc_map[ufunc]
        if method == "__call__":
            return getattr(self.mod, func_name)(*inputs, **kwargs)
    return NotImplemented

__bool__()

Converts a zero-dimensional array to a Python bool object.

Source code in src/finch/tensor/override_tensor.py

def __bool__(self):
    """
    Converts a zero-dimensional array to a Python `bool` object.
    """
    if self.ndim != 0:
        raise ValueError("Cannot convert non-scalar tensor to bool.")
    # dispatch to the scalar value's `__bool__` method
    return bool(self.item())

__complex__()

Converts a zero-dimensional array to a Python complex object.

Source code in src/finch/tensor/override_tensor.py

def __complex__(self):
    """
    Converts a zero-dimensional array to a Python `complex` object.
    """
    if self.ndim != 0:
        raise ValueError("Cannot convert non-scalar tensor to complex.")
    # dispatch to the scalar value's `__complex__` method
    return complex(self.item())

__float__()

Converts a zero-dimensional array to a Python float object.

Source code in src/finch/tensor/override_tensor.py

def __float__(self):
    """
    Converts a zero-dimensional array to a Python `float` object.
    """
    if self.ndim != 0:
        raise ValueError("Cannot convert non-scalar tensor to float.")
    # dispatch to the scalar value's `__float__` method
    return float(self.item())

__int__()

Converts a zero-dimensional array to a Python int object.

Source code in src/finch/tensor/override_tensor.py

def __int__(self):
    """
    Converts a zero-dimensional array to a Python `int` object.
    """
    if self.ndim != 0:
        raise ValueError("Cannot convert non-scalar tensor to int.")
    # dispatch to the scalar value's `__int__` method
    return int(self.item())

RandomMaskTensor

Bases: PatternTensor

A reproducible random mask with true probability p and no stored entries.

shape may be an integer for a vector, or a tuple (including ()). An omitted seed draws one random 64-bit seed from rng (a new NumPy generator by default). Reads consume no randomness; a fixed seed and coordinates give the same value even if the shape changes.

Like Finch.jl, indexed masks mix the seed, then each one-based coordinate from the first axis to the last, which is Finch.jl's order once the axes are reversed like other tensors. Scalar masks use the seed directly. The high 53 bits give a value in [0, 1) that is compared with p.

Source code in src/finch/tensor/patterns.py

class RandomMaskTensor(PatternTensor):
    """A reproducible random mask with true probability ``p`` and no stored entries.

    ``shape`` may be an integer for a vector, or a tuple (including ``()``).
    An omitted ``seed`` draws one random 64-bit seed from ``rng`` (a new NumPy
    generator by default). Reads consume no randomness; a fixed seed and
    coordinates give the same value even if the shape changes.

    Like Finch.jl, indexed masks mix the seed, then each one-based coordinate
    from the first axis to the last, which is Finch.jl's order once the axes
    are reversed like other tensors. Scalar masks use the seed directly.
    The high 53 bits give a value in ``[0, 1)`` that is compared with ``p``.
    """

    def __init__(
        self,
        shape,
        p: float,
        *,
        seed: int | None = None,
        rng: np.random.Generator | None = None,
        dtype=None,
    ):
        if not 0 <= p <= 1:
            raise ValueError("p must lie in [0, 1]")
        if isinstance(shape, int | np.integer):
            shape = (shape,)
        shape = tuple(operator.index(dim) for dim in shape)
        if any(dim < 0 for dim in shape):
            raise ValueError("shape dimensions must be nonnegative")
        if any(dim > np.iinfo(np.intp).max for dim in shape):
            raise OverflowError("shape dimensions must fit in intp")
        self._p = float(p)
        if seed is None:
            if rng is None:
                rng = np.random.default_rng()
            self._seed = int(rng.integers(0, 1 << 64, dtype=np.uint64))
        else:
            self._seed = operator.index(seed)
        if not 0 <= self._seed < 1 << 64:
            raise ValueError("seed must lie in [0, 2**64 - 1]")
        super().__init__(
            shape,
            ndim=len(shape),
            dtype=dtype,
            default_dtype=np.bool_,
            fill_value=False,
            pattern_value=True,
            p=self._p,
            seed=self._seed,
        )

    def contains(self, *idxs) -> bool:
        if not idxs:
            return _randommask_uniform(self._seed) < self._p
        state = _randommask_mix(self._seed)
        for idx in idxs:
            state = _randommask_mix(state ^ (operator.index(idx) + 1))
        return _randommask_uniform(state) < self._p

    def item(self):
        if self.ndim != 0:
            raise ValueError("Cannot convert non-scalar tensor to Python scalar.")
        return self[()].item()

RepeatTensor

Bases: PatternTensor

Map each of n indices to k repeated indices.

shape is (n, n * k), and row i holds the repeats j of i, so that the efficient row-major traversal visits each index and then its repeats. For k <= 0, all entries are false.

Source code in src/finch/tensor/patterns.py

class RepeatTensor(PatternTensor):
    """Map each of ``n`` indices to ``k`` repeated indices.

    ``shape`` is ``(n, n * k)``, and row ``i`` holds the repeats ``j`` of ``i``,
    so that the efficient row-major traversal visits each index and then its
    repeats. For ``k <= 0``, all entries are false.
    """

    def __init__(self, shape, *, k: int = 0, dtype=None):
        self._k = k
        super().__init__(shape, dtype=dtype, k=self._k)

    @classmethod
    def _level_format_properties(cls, ndim: int) -> list[FormatProperty]:
        return [Blocked((0,), (1,)), Repeated((0,), (1,))]

    def contains(self, i, j) -> bool:
        return self._k > 0 and i == j // self._k

Scalar

Bases: OverrideTensor

Source code in src/finch/tensor/scalar.py

class Scalar(OverrideTensor):
    def __init__(self, val: Any, fill_value: Any = None, device=None):
        if fill_value is None:
            fill_value = val
        self.val = val
        self._fill_value = as_fill(fill_value)
        self._device = normalize_device(device)

    @property
    def ftype(self):
        elem_t = ftype(self.val)
        return ScalarFType(elem_t, DynamicFillFType(elem_t), self._device)

    @property
    def shape(self):
        return ()

    @property
    def fill_value(self) -> Any:
        """Default value to fill the scalar."""
        return self._fill_value.value

    @property
    def device(self):
        return self._device

    def with_fill(self, fill_value: AbstractFill) -> Scalar:
        if is_dynamic(fill_value):
            return Scalar(self.val, fill_value=fill_value, device=self._device)
        if not bool(ffuncs.same(fill_value.value, self.val)):
            raise NotImplementedError(
                "a ConstantScalar's fill is its value, so a scalar of "
                f"{self.val!r} cannot take the static fill {fill_value.value!r}"
            )
        return ConstantScalar(self.val, device=self._device)

    def to_device(self, device, /, *, stream=None):
        if stream is not None:
            raise ValueError(f"stream argument is not supported; got {stream!r}")
        device = normalize_device(device)
        if device == self.device:
            return self
        return Scalar(self.val, fill_value=self._fill_value, device=device)

    @property
    def element_type(self) -> FType:
        """Data type of the scalar."""
        return self.ftype.element_type

    @property
    def shape_type(self) -> tuple:
        """Shape type of the scalar."""
        return self.ftype.shape_type

    def item(self):
        return self.val.item() if hasattr(self.val, "item") else self.val

    def __array__(self, dtype=None, copy=None):
        if copy is None:
            return np.asarray(self.val, dtype=dtype)
        return np.array(self.val, dtype=dtype, copy=copy)

    def __getitem__(self, idx):
        if idx == () or idx is Ellipsis or idx == (...,):
            return self
        raise IndexError("Too many indices for scalar tensor.")

    def __str__(self):
        return str(self.val)

    def to_numpy(self):
        return self.val

    def to_scipy(self):
        raise NotImplementedError(f"{type(self).__name__} does not support to_scipy.")

element_type property

Data type of the scalar.

fill_value property

Default value to fill the scalar.

shape_type property

Shape type of the scalar.

ScalarFType

Bases: TensorFType, ImmutableStructFType

Source code in src/finch/tensor/scalar.py

class ScalarFType(TensorFType, ImmutableStructFType):
    def __init__(
        self, _element_type: FType, _fill_type: AbstractFillFType, _device=None
    ):
        self._element_type = _element_type
        self._fill_type = _fill_type
        self._device = normalize_device(_device)

    def __eq__(self, other):
        if isinstance(other, ScalarFType):
            return (
                self._element_type == other._element_type
                and self._fill_type == other._fill_type
                and self.device == other.device
            )
        return False

    def __hash__(self):
        return hash((self._element_type, self._fill_type, self.device))

    def construct(self, shape: tuple) -> Scalar:
        if shape != ():
            raise ValueError("ScalarFType can only be called with empty shape ()")
        return self._element_type(resolve_fill(self._fill_type).value)

    def __call__(self, val: Any) -> Scalar:
        """
        Convert a tensor to this scalar tensor type.

        Args:
            val: A value to convert to this type.
        Returns:
            A Scalar instance of this type.
        """
        raise NotImplementedError(
            f"Tensor conversion not yet implemented for {type(self).__name__}"
        )

    def from_numpy(self, arr):
        return self(arr)

    @property
    def fill_type(self) -> AbstractFillFType:
        return self._fill_type

    @property
    def device(self):
        return self._device

    @property
    def element_type(self) -> FType:
        return self._element_type

    @property
    def shape_type(self):
        return ()

    @property
    def struct_name(self) -> str:
        return "Scalar"

    @property
    def struct_fields(self) -> list[tuple[str, FType]]:
        return [("val", self._element_type)]

    def from_fields(self, val):
        if is_dynamic(self._fill_type):
            # A scalar's fill defaults to its value.
            return Scalar(val, device=self._device)
        return Scalar(
            val, fill_value=resolve_fill(self._fill_type), device=self._device
        )

    def with_fill(self, fill_type: AbstractFillFType) -> ScalarFType:
        return ScalarFType(self._element_type, fill_type, self._device)

    def fisinstance(self, other):
        other_t = ftype(other)
        if is_dynamic(self._fill_type) and isinstance(other_t, ScalarFType):
            # A dynamic-fill ftype accepts any fill value of matching dtype.
            other_t = other_t.with_fill(self._fill_type)
        return other_t == self

    def lower_fill(self, ctx, obj):
        if is_static(self._fill_type):
            return asm.Literal(self._fill_type.value)
        # The dynamic fill of a scalar is its value.
        return self.lower_unwrap(ctx, obj)

    def lower_unwrap(self, ctx, obj):
        match obj:
            case ntn.Fiber():
                # A slot-bound scalar argument: read the value struct field.
                return asm.GetAttr(ctx(obj.lvl.root), asm.Literal("val"))
            case _:
                # An inline scalar value (e.g. a sparse gap read).
                return ctx(obj)

__call__(val)

Convert a tensor to this scalar tensor type.

Parameters:

Name	Type	Description	Default
val	Any	A value to convert to this type.	required

Returns: A Scalar instance of this type.

Source code in src/finch/tensor/scalar.py

def __call__(self, val: Any) -> Scalar:
    """
    Convert a tensor to this scalar tensor type.

    Args:
        val: A value to convert to this type.
    Returns:
        A Scalar instance of this type.
    """
    raise NotImplementedError(
        f"Tensor conversion not yet implemented for {type(self).__name__}"
    )

SparseListLevel dataclass

Bases: SingleDimensionLevel

A class representing sparse list level.

Source code in src/finch/tensor/level/sparse_list_level.py

@dataclass
class SparseListLevel(SingleDimensionLevel):
    """
    A class representing sparse list level.
    """

    lvl: Level
    dimension: np.integer
    ptr: Any = None
    idx: Any = None
    # These cursors must survive serialization while ptr holds per-parent counts.
    qos_fill: Any = None
    qos_stop: Any = None
    # One past the last occupied parent position; zero means none.
    prev_pos: Any = 0
    frozen: bool = True

    @property
    def shape(self) -> tuple:
        return (self.dimension, *self.lvl.shape)

    def __post_init__(self):
        if self.ptr is None:
            self.ptr = self.lvl.buffer_factory(self.lvl.position_type)(1)
        if self.idx is None:
            self.idx = self.lvl.buffer_factory(ftype(self.dimension))(0)
        p_t = self.lvl.position_type
        self.qos_fill = p_t(
            self.idx.length() if self.qos_fill is None else self.qos_fill
        )
        self.qos_stop = p_t(
            self.idx.length() if self.qos_stop is None else self.qos_stop
        )
        self.prev_pos = p_t(self.prev_pos)

    @property
    def ftype(self) -> SparseListLevelFType:
        return SparseListLevelFType(self.lvl.ftype, ftype(self.dimension))

    @property
    def val(self) -> Any:
        return self.lvl.val

    def __str__(self):
        return f"SparseListLevel(lvl={self.lvl}, dim={self.dimension})"

SparseListLevelFType dataclass

Bases: SingleDimensionLevelFType, MutableStructFType

Source code in src/finch/tensor/level/sparse_list_level.py

@dataclass(unsafe_hash=True)
class SparseListLevelFType(SingleDimensionLevelFType, MutableStructFType):
    _lvl_t: LevelFType
    dimension_type: ftypes.FDTypeInteger = ftypes.intp

    def __post_init__(self):
        self.dimension_type = ftype(self.dimension_type)

    @property
    def struct_name(self):
        return "SparseListLevelFType"

    @property
    def p_t(self):
        return self.position_type

    @property
    def struct_fields(self):
        return [
            ("lvl", self.lvl_t),
            ("dimension", self.dimension_type),
            ("ptr", self.ptr_type),
            ("idx", self.idx_type),
            ("qos_fill", self.position_type),
            ("qos_stop", self.position_type),
            ("prev_pos", self.position_type),
            ("frozen", ftypes.bool_),
        ]

    def __str__(self):
        return f"SparseListLevelFType({self.lvl_t})"

    @property
    def ndim(self):
        return 1 + self.lvl_t.ndim

    @property
    def fill_type(self):
        return self.lvl_t.fill_type

    @property
    def element_type(self):
        """
        Returns the type of elements stored in the fibers.
        """
        return self.lvl_t.element_type

    @property
    def shape_type(self):
        """
        Returns the type of the shape of the fibers.
        """
        return (self.dimension_type, *self.lvl_t.shape_type)

    @property
    def position_type(self):
        """
        Returns the type of positions within the levels.
        """
        return self.lvl_t.position_type

    @property
    def buffer_type(self):
        return self.lvl_t.buffer_type

    @property
    def buffer_factory(self):
        """
        Returns the ftype of the buffer used for the fibers.
        """
        return self.lvl_t.buffer_factory

    @property
    def ptr_type(self):
        return self.buffer_factory(self.position_type)

    @property
    def idx_type(self):
        return self.buffer_factory(self.dimension_type)

    def level_cost(self, fields, stats, stats_factory, num_pos, lvl) -> float:
        pos_size = np_dtype(self.position_type).itemsize
        size_ptr = (num_pos + 1) * pos_size
        reduce_fields = fields[lvl + 1 :]
        if reduce_fields:
            reduced_stats = stats_factory.aggregate(
                ffuncs.or_, False, reduce_fields, stats_factory.support(stats)
            )
        else:
            reduced_stats = stats
        nnz_prefix = reduced_stats.estimate_non_fill_values()
        size_idx = nnz_prefix * pos_size

        return (
            size_ptr
            + size_idx
            + self.lvl_t.level_cost(fields, stats, stats_factory, nnz_prefix, lvl + 1)
        )

    def construct(self, shape: tuple[Any, ...], *, pos: int) -> "SparseListLevel":
        """
        Creates an instance of SparseListLevel.

        Args:
            shape: The shape to be used for the level. (mandatory)
        Returns:
            An instance of DenseLevel.
        """
        lvl = self.lvl_t.construct(shape=shape[1:], pos=0)
        return SparseListLevel(
            lvl,
            self.dimension_type(shape[0]),
            self.ptr_type(int(pos) + 1),
            self.idx_type(0),
        )

    def __call__(self, val: Any) -> "SparseListLevel":
        """
        Convert a level to this sparse list level type.

        Args:
            val: A value to convert to this type.
        Returns:
            A SparseListLevel instance of this type.
        """
        raise NotImplementedError(
            f"Level conversion not yet implemented for {type(self).__name__}"
        )

    @property
    def lvl_t(self):
        return self._lvl_t

    def level_format_properties(self, n):
        return self.lvl_t.level_format_properties(n + 1)

    def level_lower_dim(self, ctx, obj, r):
        if r == 0:
            return asm.GetAttr(obj, asm.Literal("dimension"))
        return self.lvl_t.level_lower_dim(
            ctx, asm.GetAttr(obj, asm.Literal("lvl")), r - 1
        )

    def level_lower_declare(self, ctx, tns, init, op, shape, pos):
        p_t = self.ptr_type.length_type
        p = asm.Variable(ctx.freshen("p"), p_t)
        stop = asm.Variable(ctx.freshen("pos_stop"), p_t)
        to_size = asm.Literal(ffuncs.astype(p_t))
        zero = asm.Literal(self.position_type(0))
        expr = """finch
        stop = to_size(pos) + 1
        resize(tns.ptr, stop)
        for (p in 0:stop)
            tns.ptr[p] = zero
        end
        tns.qos_fill = zero
        tns.qos_stop = zero
        tns.prev_pos = zero
        tns.frozen = false
        """
        ctx.exec(parse_assembly(expr, locals(), position_type=p_t))
        return self.lvl_t.level_lower_declare(
            ctx, asm.GetAttr(tns, asm.Literal("lvl")), init, op, shape, zero
        )

    def level_lower_thaw(self, ctx, lvl, op, pos):
        p_t = self.ptr_type.length_type
        pos_stop = asm.Variable(ctx.freshen("pos_stop"), p_t)
        p = asm.Variable(ctx.freshen("p"), p_t)
        to_size = asm.Literal(ffuncs.astype(p_t))
        to_pos = asm.Literal(ffuncs.astype(self.position_type))
        zero = asm.Literal(self.position_type(0))
        expr = """finch
        pos_stop = to_size(pos)
        if (lvl.frozen)
            lvl.qos_fill = lvl.ptr[pos_stop]
            lvl.qos_stop = lvl.qos_fill
            lvl.prev_pos = zero
            p = pos_stop
            // Difference backwards while preceding entries are still prefix sums.
            while (p > 0)
                lvl.ptr[p] -= lvl.ptr[p - 1]
                if (lvl.ptr[p] != zero)
                    if (lvl.prev_pos == zero)
                        lvl.prev_pos = to_pos(p)
                    end
                end
                p -= 1
            end
            lvl.frozen = false
        end
        """
        ctx.exec(parse_assembly(expr, locals(), position_type=p_t))
        return self.lvl_t.level_lower_thaw(
            ctx,
            asm.GetAttr(lvl, asm.Literal("lvl")),
            op,
            asm.GetAttr(lvl, asm.Literal("qos_fill")),
        )

    def level_lower_freeze(self, ctx, lvl, op, pos):
        p_t = self.ptr_type.length_type
        pos_stop = asm.Variable(ctx.freshen("pos_stop"), p_t)
        qos_stop = asm.Variable(ctx.freshen("qos_stop"), self.position_type)
        p = asm.Variable(ctx.freshen("p"), p_t)
        to_size = asm.Literal(ffuncs.astype(p_t))
        to_idx_size = asm.Literal(ffuncs.astype(self.idx_type.length_type))
        expr = """finch
        pos_stop = to_size(pos)
        resize(lvl.ptr, pos_stop + 1)
        if (lvl.frozen == false)
            for (p in 0:pos_stop)
                lvl.ptr[p + 1] += lvl.ptr[p]
            end
        end
        qos_stop = lvl.ptr[pos_stop]
        resize(lvl.idx, to_idx_size(qos_stop))
        lvl.qos_fill = qos_stop
        lvl.qos_stop = qos_stop
        lvl.frozen = true
        """
        ctx.exec(parse_assembly(expr, locals(), position_type=p_t))
        return self.lvl_t.level_lower_freeze(
            ctx,
            asm.GetAttr(lvl, asm.Literal("lvl")),
            op,
            asm.GetAttr(lvl, asm.Literal("qos_fill")),
        )

    def level_lower_increment(self, ctx, obj, op, val, pos):
        raise NotImplementedError(
            "SparseListLevelFType does not support level_lower_increment."
        )

    def level_lower_assemble(self, ctx, lvl, start, stop):
        p_t = self.ptr_type.length_type
        p = asm.Variable(ctx.freshen("p"), p_t)
        p_start = asm.Variable(ctx.freshen("p_start"), p_t)
        p_stop = asm.Variable(ctx.freshen("p_stop"), p_t)
        to_size = asm.Literal(ffuncs.astype(p_t))
        length = asm.Length(asm.GetAttr(lvl, asm.Literal("ptr")))
        zero = asm.Literal(self.position_type(0))
        expr = """finch
        p_start = to_size(start) + 1
        p_stop = to_size(stop) + 1
        if (length < p_stop)
            resize(lvl.ptr, p_stop)
        end
        for (p in p_start:p_stop)
            lvl.ptr[p] = zero
        end
        """
        ctx.exec(parse_assembly(expr, locals(), position_type=p_t))

    def level_lower_unwrap(self, ctx, obj, pos):
        raise NotImplementedError(
            "SparseListLevelFType does not support level_lower_unwrap."
        )

    def level_unfurl(
        self, ctx, fiber: ntn.Fiber, ext, mode: ntn.AccessMode, proto, pos
    ):
        match mode:
            case ntn.Update():
                return self.level_unfurl_update(ctx, fiber, ext, mode, proto, pos)
        tns = fiber
        level = tns.lvl
        lvl_asm = ctx(level)
        ptr_s = asm.GetAttr(lvl_asm, asm.Literal("ptr"))
        idx_s = asm.GetAttr(lvl_asm, asm.Literal("idx"))

        q = asm.Variable(ctx.freshen("q"), self.position_type)
        q_stop = asm.Variable(ctx.freshen("q_stop"), self.position_type)
        i_stop = asm.Variable(ctx.freshen("i_stop"), self.position_type)
        i_last = asm.Variable(ctx.freshen("i_last"), self.position_type)
        fill = (
            ntn.Value(self.lower_fill(lvl_asm), self.element_type)
            if is_dynamic(self.fill_type)
            else ntn.Literal(self.fill_type.value)
        )
        full = ntn.Full(
            fill,
            tuple(
                ntn.Value(self.level_lower_dim(ctx, lvl_asm, r), self.shape_type[r])
                for r in range(1, self.ndim)
            ),
        )
        tmp_locals = locals()

        def thunk_preamble(ctx, idx):
            expr = """finch
            q = ptr_s[pos]
            q_stop = ptr_s[pos + 1]
            i_stop = 1
            i_last = 0
            if (q < q_stop)
                i_stop = idx_s[q]
                i_last = idx_s[q_stop - 1]
            end
            """
            return parse_assembly(expr, tmp_locals, position_type=self.position_type)

        def seek_fn(ctx, ext):
            start = ctx.ctx(ext.get_start())

            code = f"""finch
            if (idx_s[q] < {start})
              q = scansearch(idx_s, {start}, q, q_stop - 1)
            end
            """

            return parse_assembly(
                code,
                tmp_locals | asm.get_vars_in_expr(start),
                position_type=self.position_type,
            )

        def chunk_tail_fn(ctx, idx):
            pos_2 = asm.Variable(
                ctx.freshen(idx, f"_pos_{self.ndim - 1}"), self.position_type
            )
            ctx.exec(asm.Assign(pos_2, q))
            return lplt.Run(
                ntn.Fiber(
                    ntn.Child(level),
                    ntn.Value(pos_2, self.position_type),
                    (*tns.idxs, idx),
                )
            )

        return lplt.Thunk(
            preamble=thunk_preamble,
            body=lambda ctx, ext: lplt.Sequence(
                head=lambda ctx, idx: lplt.Stepper(
                    preamble=lambda ctx: asm.IfElse(
                        asm.Call(asm.L(ffuncs.lt), (q, q_stop)),
                        asm.Block((asm.Assign(i_stop, asm.Load(idx_s, q)),)),
                        asm.Block(
                            (
                                asm.Assign(
                                    i_stop,
                                    asm.GetAttr(lvl_asm, asm.Literal("dimension")),
                                ),
                            )
                        ),
                    ),
                    stop=lambda ctx: ntn.Variable(i_stop.name, self.position_type),
                    chunk=lplt.Sequence(
                        head=lambda ctx, idx: lplt.Run(full),
                        split=lambda ctx, ext: ntn.Variable(
                            i_stop.name, self.position_type
                        ),
                        tail=chunk_tail_fn,
                    ),
                    next=lambda ctx: asm.Block(
                        (
                            asm.Assign(
                                q,
                                asm.Call(asm.L(ffuncs.add), (q, asm.L(self.p_t(1)))),
                            ),
                        )
                    ),
                    seek=seek_fn,
                ),
                split=lambda ctx, idx: ntn.Call(
                    ntn.L(ffuncs.add),
                    (ntn.Variable(i_last.name, self.position_type), ext.get_unit()),
                ),
                tail=lambda ctx, idx: lplt.Run(full),
            ),
        )

    def level_unfurl_update(self, ctx, fiber, ext, mode, proto, pos):
        lvl = ctx(fiber.lvl)
        child = asm.GetAttr(lvl, asm.Literal("lvl"))
        p_t = self.position_type
        qos = asm.Variable(ctx.freshen("qos"), p_t)
        dirty = ntn.Variable(ctx.freshen("dirty"), ftypes.bool_)
        dirty_asm = ctx(dirty)
        to_ptr_size = asm.Literal(ffuncs.astype(self.ptr_type.length_type))
        to_idx_size = asm.Literal(ffuncs.astype(self.idx_type.length_type))
        to_index = asm.Literal(ffuncs.astype(self.dimension_type))
        to_pos = asm.Literal(ffuncs.astype(p_t))
        bindings = locals()

        def preamble(ctx, idx):
            body = parse_assembly(
                """finch
            qos = lvl.qos_fill
            """,
                bindings,
                position_type=p_t,
            )
            if ctx.mode.safe or ctx.mode.debug:
                return asm.Block(
                    (
                        *body.bodies,
                        asm.Assert(
                            asm.Call(
                                asm.Literal(ffuncs.not_),
                                (asm.GetAttr(lvl, asm.Literal("frozen")),),
                            )
                        ),
                        asm.Assert(
                            asm.Call(
                                asm.Literal(ffuncs.le),
                                (asm.GetAttr(lvl, asm.Literal("prev_pos")), pos),
                            )
                        ),
                    )
                )
            return body

        def lookup(ctx, idx):
            idx_asm = ctx.ctx(idx)

            def prepare(ctx, _):
                grow = ctx.ctx.block()
                grow.exec(
                    parse_assembly(
                        """finch
                if (lvl.qos_stop == 0)
                    lvl.qos_stop = 1
                else
                    lvl.qos_stop += lvl.qos_stop
                end
                resize(lvl.idx, to_idx_size(lvl.qos_stop))
                """,
                        bindings,
                        position_type=p_t,
                    )
                )
                self.lvl_t.level_lower_assemble(
                    grow, child, qos, asm.GetAttr(lvl, asm.Literal("qos_stop"))
                )
                return asm.Block(
                    (
                        asm.If(
                            asm.Call(
                                asm.Literal(ffuncs.ge),
                                (qos, asm.GetAttr(lvl, asm.Literal("qos_stop"))),
                            ),
                            asm.Block(tuple(grow.emit())),
                        ),
                        asm.Assign(dirty_asm, asm.Literal(False)),
                    )
                )

            def finish(ctx, _):
                written = parse_assembly(
                    """finch
                lvl.idx[to_idx_size(qos)] = to_index(idx_asm)
                qos += 1
                lvl.prev_pos = to_pos(pos + 1)
                """,
                    bindings | {"idx_asm": idx_asm},
                    position_type=p_t,
                )
                match fiber:
                    case ntn.HollowFiber(dirty=parent_dirty):
                        written = asm.Block(
                            (
                                *written.bodies,
                                asm.Assign(ctx.ctx(parent_dirty), asm.Literal(True)),
                            )
                        )
                return asm.If(dirty_asm, written)

            return lplt.Thunk(
                preamble=prepare,
                body=lambda ctx, ext: lplt.Run(
                    ntn.HollowFiber(
                        ntn.Child(fiber.lvl),
                        ntn.Value(qos, p_t),
                        (*fiber.idxs, idx),
                        dirty=dirty,
                    )
                ),
                epilogue=finish,
            )

        def epilogue(ctx, idx):
            return parse_assembly(
                """finch
            lvl.ptr[to_ptr_size(pos + 1)] += qos - lvl.qos_fill
            lvl.qos_fill = qos
            """,
                bindings,
                position_type=p_t,
            )

        return lplt.Thunk(
            preamble=preamble,
            body=lambda ctx, ext: lplt.Lookup(lookup),
            epilogue=epilogue,
        )

    def from_fields(
        self,
        lvl,
        dimension,
        ptr,
        idx,
        qos_fill=None,
        qos_stop=None,
        prev_pos=0,
        frozen=True,
    ) -> "SparseListLevel":
        return SparseListLevel(
            lvl, dimension, ptr, idx, qos_fill, qos_stop, prev_pos, frozen
        )

buffer_factory property

Returns the ftype of the buffer used for the fibers.

element_type property

Returns the type of elements stored in the fibers.

position_type property

Returns the type of positions within the levels.

shape_type property

Returns the type of the shape of the fibers.

__call__(val)

Convert a level to this sparse list level type.

Parameters:

Name	Type	Description	Default
val	Any	A value to convert to this type.	required

Returns: A SparseListLevel instance of this type.

Source code in src/finch/tensor/level/sparse_list_level.py

def __call__(self, val: Any) -> "SparseListLevel":
    """
    Convert a level to this sparse list level type.

    Args:
        val: A value to convert to this type.
    Returns:
        A SparseListLevel instance of this type.
    """
    raise NotImplementedError(
        f"Level conversion not yet implemented for {type(self).__name__}"
    )

construct(shape, *, pos)

Creates an instance of SparseListLevel.

Parameters:

Name	Type	Description	Default
shape	tuple[Any, ...]	The shape to be used for the level. (mandatory)	required

Returns: An instance of DenseLevel.

Source code in src/finch/tensor/level/sparse_list_level.py

def construct(self, shape: tuple[Any, ...], *, pos: int) -> "SparseListLevel":
    """
    Creates an instance of SparseListLevel.

    Args:
        shape: The shape to be used for the level. (mandatory)
    Returns:
        An instance of DenseLevel.
    """
    lvl = self.lvl_t.construct(shape=shape[1:], pos=0)
    return SparseListLevel(
        lvl,
        self.dimension_type(shape[0]),
        self.ptr_type(int(pos) + 1),
        self.idx_type(0),
    )

SplitMaskTensor

Bases: PatternTensor

Partition n indices into p contiguous regions of nearly equal size.

shape is (p, n), and row j covers n*j//p <= i < n*(j+1)//p, so that the efficient row-major traversal visits each region and then its indices.

Source code in src/finch/tensor/patterns.py

class SplitMaskTensor(PatternTensor):
    """Partition ``n`` indices into ``p`` contiguous regions of nearly equal size.

    ``shape`` is ``(p, n)``, and row ``j`` covers ``n*j//p <= i < n*(j+1)//p``,
    so that the efficient row-major traversal visits each region and then its
    indices.
    """

    def __init__(self, shape, *, dtype=None):
        super().__init__(
            shape,
            dtype=dtype,
            default_dtype=np.bool_,
            fill_value=False,
            pattern_value=True,
        )
        p, n = map(operator.index, self.shape)
        if n < 0:
            raise ValueError("n must be nonnegative")
        if p <= 0:
            raise ValueError("shape[0] must be positive")

    def contains(self, j, i) -> bool:
        p, n = self.shape
        return n * j // p <= i < n * (j + 1) // p

Tensor

Bases: FTyped[TensorFType[FT]], ABC

Abstract base class for tensor-like data structures. Tensors are multi-dimensional arrays that can be used to represent data in various formats. They support operations such as indexing, slicing, and reshaping, and can be used in mathematical computations. This class provides the basic interface for tensors to be used with lazy ops in Finch, though more advanced interfaces may be required for different backends.

Source code in src/finch/algebra/tensor.py

class Tensor(FTyped[TensorFType[FT]], ABC):
    """
    Abstract base class for tensor-like data structures. Tensors are
    multi-dimensional arrays that can be used to represent data in various
    formats. They support operations such as indexing, slicing, and reshaping,
    and can be used in mathematical computations. This class provides the basic
    interface for tensors to be used with lazy ops in Finch, though more
    advanced interfaces may be required for different backends.
    """

    @property
    def ndim(self) -> int:
        """Number of dimensions of the tensor."""
        return int(self.ftype.ndim)

    @property
    def size(self):
        size = 1
        for dim in self.shape:
            size *= int(dim)
        return size

    @property
    @abstractmethod
    def ftype(self) -> TensorFType[FT]:
        """FType of the tensor, which may include metadata about the tensor."""
        ...

    @property
    def fill_value(self) -> Any:
        """The fill value for the tensor.  The fill value is the
        default value for a tensor when it is created with a given shape and dtype,
        as well as the background value for sparse tensors.
        Note: this is the raw value NOT the AbstractFill. A tensor whose fill is
        dynamic holds its own fill value.
        """
        fill_type = self.ftype.fill_type
        if is_static(fill_type):
            return fill_type.value
        raise NotImplementedError(
            f"{type(self).__name__} must hold its dynamic fill value"
        )

    @property
    def fill(self) -> AbstractFill:
        """
        The fill of the tensor: its fill value, and whether kernels may
        specialize on it. Its ftype is the fill of the tensor's format.
        """
        fill_type = self.ftype.fill_type
        if is_static(fill_type):
            return StaticFill(fill_type.value)
        return DynamicFill(self.fill_value, fill_type.element_type)

    def with_fill(self, fill_value: AbstractFill) -> Tensor:
        """
        An equivalent tensor carrying `fill_value` as its fill.
        """
        raise NotImplementedError(
            f"{type(self).__name__} cannot replace its fill value"
        )

    @property
    def element_type(self) -> FT:
        """The element type of the tensor.  The element type is the scalar type of
        the elements in a tensor, which may be different from the data type of the
        tensor.
        """
        return self.ftype.element_type

    @property
    def dtype(self):
        return self.element_type

    @property
    def device(self):
        return getattr(self, "_device", self.ftype.device)

    @property
    def device_type(self) -> FType:
        return self.ftype.device_type

    def to_device(self, device, /, *, stream=None):
        if stream is not None:
            raise ValueError(f"stream argument is not supported; got {stream!r}")
        device = normalize_device(device)
        if device == self.device:
            return self
        raise ValueError(f"device argument is not supported; got {device!r}")

    @property
    def shape_type(self) -> tuple[FType, ...]:
        """Shape type of the tensor. The shape type is a tuple of the index
        types in the tensor. It's the type of each element in tns.shape. It
        should be an actual tuple, rather than a tupleftype, so that it can hold
        e.g. dtypes, formats, or types, and so that we can easily index it."""
        return self.ftype.shape_type

    @property
    @abstractmethod
    def shape(self) -> tuple[np.intp, ...]:
        """Shape of the tensor."""
        ...

    @abstractmethod
    def item(self):
        """
        Return the Python scalar stored in a zero-dimensional tensor.

        Indexing tensor objects must keep returning tensor-like values, including for
        zero-dimensional results. `item()` is the explicit escape hatch for code that
        needs the underlying Python scalar value. Implementations should raise
        `ValueError` when called on a non-scalar tensor.
        """
        ...

    @abstractmethod
    def to_numpy(self):
        """Convert the tensor to a NumPy representation."""
        ...

    @abstractmethod
    def to_scipy(self):
        """Convert the tensor to a SciPy sparse representation."""
        ...

element_type property

The element type of the tensor. The element type is the scalar type of the elements in a tensor, which may be different from the data type of the tensor.

fill property

The fill of the tensor: its fill value, and whether kernels may specialize on it. Its ftype is the fill of the tensor's format.

fill_value property

The fill value for the tensor. The fill value is the default value for a tensor when it is created with a given shape and dtype, as well as the background value for sparse tensors. Note: this is the raw value NOT the AbstractFill. A tensor whose fill is dynamic holds its own fill value.

ftype abstractmethod property

FType of the tensor, which may include metadata about the tensor.

ndim property

Number of dimensions of the tensor.

shape abstractmethod property

Shape of the tensor.

shape_type property

Shape type of the tensor. The shape type is a tuple of the index types in the tensor. It's the type of each element in tns.shape. It should be an actual tuple, rather than a tupleftype, so that it can hold e.g. dtypes, formats, or types, and so that we can easily index it.

item() abstractmethod

Return the Python scalar stored in a zero-dimensional tensor.

Indexing tensor objects must keep returning tensor-like values, including for zero-dimensional results. item() is the explicit escape hatch for code that needs the underlying Python scalar value. Implementations should raise ValueError when called on a non-scalar tensor.

Source code in src/finch/algebra/tensor.py

@abstractmethod
def item(self):
    """
    Return the Python scalar stored in a zero-dimensional tensor.

    Indexing tensor objects must keep returning tensor-like values, including for
    zero-dimensional results. `item()` is the explicit escape hatch for code that
    needs the underlying Python scalar value. Implementations should raise
    `ValueError` when called on a non-scalar tensor.
    """
    ...

to_numpy() abstractmethod

Convert the tensor to a NumPy representation.

Source code in src/finch/algebra/tensor.py

@abstractmethod
def to_numpy(self):
    """Convert the tensor to a NumPy representation."""
    ...

to_scipy() abstractmethod

Convert the tensor to a SciPy sparse representation.

Source code in src/finch/algebra/tensor.py

@abstractmethod
def to_scipy(self):
    """Convert the tensor to a SciPy sparse representation."""
    ...

with_fill(fill_value)

An equivalent tensor carrying fill_value as its fill.

Source code in src/finch/algebra/tensor.py

def with_fill(self, fill_value: AbstractFill) -> Tensor:
    """
    An equivalent tensor carrying `fill_value` as its fill.
    """
    raise NotImplementedError(
        f"{type(self).__name__} cannot replace its fill value"
    )

TensorFType

Bases: FType, ABC, Generic[FT]

Source code in src/finch/algebra/tensor.py

class TensorFType(FType, ABC, Generic[FT]):
    @property
    def ndim(self) -> int:
        """Number of dimensions of the tensor."""
        return len(self.shape_type)

    @property
    def shape(self) -> tuple[int, ...]:
        raise TypeError("TensorFType does not have a shape.")

    @property
    def size(self) -> int:
        return functools.reduce(lambda x, y: x * int(y), self.shape, 1)

    @property
    @abstractmethod
    def fill_type(self) -> AbstractFillFType:
        """The fill of the format, either static or dynamic."""
        ...

    def with_fill(self, fill_type: AbstractFillFType) -> TensorFType:
        """
        This ftype carrying `fill_type` as its fill.
        """
        raise NotImplementedError(
            f"{type(self).__name__} cannot replace its fill value"
        )

    @property
    @abstractmethod
    def element_type(self) -> FT:
        """Data type of the tensor elements."""
        ...

    @property
    def dtype(self):
        return self.element_type

    @property
    def device(self):
        return serial()

    @property
    def device_type(self) -> FType:
        return self.device.ftype

    @property
    @abstractmethod
    def shape_type(self) -> tuple[FType, ...]:
        """Shape type of the tensor. The shape type is a tuple of the index
        types in the tensor. It's the type of each element in tns.shape. It
        should be an actual tuple, rather than a tuple type, so that it can hold
        e.g. dtypes, formats, or types, and so that we can easily index it."""
        ...

    @abstractmethod
    def construct(self, shape: tuple) -> Tensor:
        """
        Create a tensor instance with the given shape.

        Args:
            shape: The shape of the tensor to create.
        Returns:
            A tensor instance with the specified shape.
        """
        ...

    @abstractmethod
    def from_numpy(self, arr: np.ndarray) -> Tensor: ...

element_type abstractmethod property

Data type of the tensor elements.

fill_type abstractmethod property

The fill of the format, either static or dynamic.

ndim property

Number of dimensions of the tensor.

shape_type abstractmethod property

Shape type of the tensor. The shape type is a tuple of the index types in the tensor. It's the type of each element in tns.shape. It should be an actual tuple, rather than a tuple type, so that it can hold e.g. dtypes, formats, or types, and so that we can easily index it.

construct(shape) abstractmethod

Create a tensor instance with the given shape.

Parameters:

Name	Type	Description	Default
shape	tuple	The shape of the tensor to create.	required

Returns: A tensor instance with the specified shape.

Source code in src/finch/algebra/tensor.py

@abstractmethod
def construct(self, shape: tuple) -> Tensor:
    """
    Create a tensor instance with the given shape.

    Args:
        shape: The shape of the tensor to create.
    Returns:
        A tensor instance with the specified shape.
    """
    ...

with_fill(fill_type)

This ftype carrying fill_type as its fill.

Source code in src/finch/algebra/tensor.py

def with_fill(self, fill_type: AbstractFillFType) -> TensorFType:
    """
    This ftype carrying `fill_type` as its fill.
    """
    raise NotImplementedError(
        f"{type(self).__name__} cannot replace its fill value"
    )

asarray(obj, /, *, dtype=None, device=None, copy=None, format=None)

asarray(
    obj: np.ndarray,
    /,
    *,
    dtype=None,
    device=None,
    copy=None,
    format: None = None,
) -> BufferizedNDArray

asarray(
    obj: sps.spmatrix | sps.sparray,
    /,
    *,
    dtype=None,
    device=None,
    copy=None,
    format: None = None,
) -> FiberTensor

asarray(
    obj: Any,
    /,
    *,
    dtype=None,
    device=None,
    copy=None,
    format: FiberTensorFType,
) -> FiberTensor

asarray(
    obj: np.generic | complex,
    /,
    *,
    dtype=None,
    device=None,
    copy=None,
    format: None = None,
) -> Scalar

asarray(
    obj: Any,
    /,
    *,
    dtype=None,
    device=None,
    copy=None,
    format: TensorFType | None = None,
) -> Tensor

Convert given argument and return wrapper type instance. If input argument is already array type, return unchanged. https://data-apis.org/array-api/latest/API_specification/generated/array_api.asarray.html

Source code in src/finch/interface/lazy.py

def asarray(
    obj: Any,
    /,
    *,
    dtype=None,
    device=None,
    copy=None,
    format: TensorFType | None = None,
) -> Tensor:
    """
    Convert given argument and return wrapper type instance.
    If input argument is already array type, return unchanged.
    https://data-apis.org/array-api/latest/API_specification/generated/array_api.asarray.html
    """
    explicit_device = device is not None
    device = normalize_device(device)

    if format is None:
        if isinstance(obj, Scalar):
            if copy is True or (explicit_device and obj.device != device):
                return Scalar(obj.val, fill_value=obj.fill_value, device=device)
            return obj
        if isinstance(obj, BufferizedNDArray):
            if copy is True:
                return BufferizedNDArray.from_numpy(
                    obj.to_numpy().copy(), fill_value=obj.fill_value, device=device
                )
            if explicit_device and obj.device != device:
                return obj.to_device(device)
            return obj
        if isinstance(obj, Tensor) and copy is not True:
            if explicit_device and obj.device != device:
                return obj.to_device(device)
            return obj
        if isinstance(obj, np.ndarray):
            if copy is True:
                obj = obj.copy()
            return BufferizedNDArray.from_numpy(obj, device=device)
        if sps.issparse(obj):
            if obj.format == "csc":  # CSC has no native Finch layout yet
                if dtype is not None or copy is True:
                    raise NotImplementedError(
                        "asarray cannot honor dtype or copy for CSC input; "
                        "convert to CSR or COO first."
                    )
                return obj
            if copy is False and not (
                obj.format in ("coo", "csr") and obj.has_canonical_format
            ):
                raise ValueError(
                    "Unable to avoid copy while creating an array as requested."
                )
            if obj.format not in ("coo", "csr"):
                obj = obj.asformat("coo")

            if dtype is not None:
                if copy is False:
                    obj = obj.astype(dtype, copy=False)
                else:
                    obj = obj.astype(dtype, copy=True)
            elif copy is True:
                obj = obj.copy()

            if not obj.has_canonical_format:
                obj = obj.copy()
                obj.sum_duplicates()

            match obj.format:
                case "csr":
                    return FiberTensor.from_scipy_csr(obj, device=device)
                case "coo":
                    return FiberTensor.from_scipy_coo(obj, device=device)

        if np.isscalar(obj) or obj is None:
            if dtype is not None:
                obj = ftype(dtype)(obj)
            elif isinstance(obj, bool | int | float | complex):
                obj = np.asarray(obj)[()]
            return Scalar(obj, device=device)
        try:
            np_arr = np.asarray(obj)
            if np_arr.dtype != object:
                if dtype is not None:
                    np_arr = np_arr.astype(np_dtype(ftype(dtype)))
                elif copy is True:
                    np_arr = np_arr.copy()
                return BufferizedNDArray.from_numpy(np_arr, device=device)
        except (TypeError, ValueError):
            pass
        return obj

    out = format.from_numpy(obj) if isinstance(obj, np.ndarray) else format(obj)
    return out.to_device(device) if explicit_device else out

broadcast_arrays(*args)

Broadcasts one or more arrays against one another.

Parameters

*args: array an arbitrary number of to-be broadcasted arrays.

Returns

out: List[array] a list of broadcasted arrays. Each array has the same shape. Element types are preserved.

Source code in src/finch/interface/eager.py

def broadcast_arrays(*args):
    """
    Broadcasts one or more arrays against one another.

    Parameters
    ----------
    *args: array
        an arbitrary number of to-be broadcasted arrays.

    Returns
    -------
    out: List[array]
        a list of broadcasted arrays. Each array has the same shape.
        Element types are preserved.
    """
    if builtins.any(isinstance(arg, lazy.LazyTensor) for arg in args):
        return lazy.broadcast_arrays(*args)
    # compute can take in a list of LazyTensors
    return compute(lazy.broadcast_arrays(*args))

broadcast_to(x, /, shape)

Broadcasts an array to a new shape.

Parameters

x: array The input tensor to be broadcasted. shape: Sequence[int] The target shape to which the input tensor should be broadcasted.

Returns

out: array A tensor with the same data as x, but with the specified shape.

Source code in src/finch/interface/eager.py

def broadcast_to(x, /, shape: Sequence[int]):
    """
    Broadcasts an array to a new shape.

    Parameters
    ----------
    x: array
        The input tensor to be broadcasted.
    shape: Sequence[int]
        The target shape to which the input tensor should be broadcasted.

    Returns
    -------
    out: array
        A tensor with the same data as `x`, but with the specified shape.
    """
    shape = tuple(shape)  # Ensure shape is a tuple for consistency
    if isinstance(x, lazy.LazyTensor):
        return lazy.broadcast_to(x, shape=shape)
    return compute(lazy.broadcast_to(x, shape=shape))

compute(arg, ctx=None)

Executes a fused operation represented by LazyTensors. This function evaluates the entire operation in an optimized manner using the provided scheduler.

• arg: A lazy tensor or a tuple of lazy tensors representing the fused operation to be computed.
• ctx: The scheduler to use for computation. Defaults to the result of get_default_scheduler().

Returns: - A tensor or a list of tensors computed by the fused operation.

Source code in src/finch/interface/fuse.py

def compute(arg, ctx=None):
    """
    Executes a fused operation represented by LazyTensors. This function evaluates the
    entire operation in an optimized manner using the provided scheduler.

    Parameters:
    - arg: A lazy tensor or a tuple of lazy tensors representing the fused operation to
      be computed.
    - ctx: The scheduler to use for computation. Defaults to the result of
      `get_default_scheduler()`.

    Returns:
    - A tensor or a list of tensors computed by the fused operation.
    """
    if ctx is None:
        ctx = get_default_scheduler()

    args = arg if isinstance(arg, tuple) else (arg,)
    outputs = [None for _ in args]
    lazy_args = []
    lazy_arg_idxs = []
    for i, arg_i in enumerate(args):
        if not isinstance(arg_i, LazyTensor):
            outputs[i] = arg_i
        else:
            lazy_args.append(arg_i)
            lazy_arg_idxs.append(i)

    if lazy_args:
        vars = tuple(Alias(gensym("A")) for _ in lazy_args)
        ctx_2 = lazy_args[0].ctx.join(*[x.ctx for x in lazy_args[1:]])
        bodies = []
        for arg_i, var in zip(lazy_args, vars, strict=True):
            idxs = tuple(Field(gensym("i")) for _ in range(len(arg_i.shape)))
            bodies.append(Query(Table(var, idxs), Table(arg_i.data, idxs)))
        prgm = Plan((*ctx_2.trace(), *bodies, Produces(vars)))
        res = ctx(prgm)
        for lazy_idx, out_idx in enumerate(lazy_arg_idxs):
            device = lazy_args[lazy_idx].device
            if (
                len(res[lazy_idx].shape) == 0
            ):  # if the result is a scalar, extract the value and turn it into a
                # finch `Scalar`
                outputs[out_idx] = asarray(res[lazy_idx][()], device=device)
            else:
                outputs[out_idx] = asarray(res[lazy_idx], device=device)

    return tuple(outputs) if isinstance(arg, tuple) else outputs[0]

defer(arr)

defer(arr: tuple[Any, ...]) -> tuple[Any, ...]

defer(arr: Any) -> LazyTensor

• defer(arr) -> LazyTensor: Converts an array into a LazyTensor. If the input is already a LazyTensor, it is returned as-is. Otherwise, it creates a LazyTensor representation of the input array.

Parameters: - arr: The input array to be converted into a LazyTensor.

Returns: - LazyTensor: A lazy representation of the input array.

Source code in src/finch/interface/lazy.py

def defer(arr: Any) -> LazyTensor | tuple[Any, ...]:
    """
    - defer(arr) -> LazyTensor:
    Converts an array into a LazyTensor. If the input is already a LazyTensor, it is
    returned as-is.
    Otherwise, it creates a LazyTensor representation of the input array.

    Parameters:
    - arr: The input array to be converted into a LazyTensor.

    Returns:
    - LazyTensor: A lazy representation of the input array.
    """
    if isinstance(arr, tuple):
        return tuple(defer(arr_i) for arr_i in arr)

    if isinstance(arr, LazyTensor):
        return arr
    if _is_numeric_constant(arr):
        arr = ConstantScalar(arr) if is_specializable_value(arr) else Scalar(arr)
    else:
        arr = asarray(arr)
    tns = Alias(gensym("A"))
    idxs = tuple(Field(gensym("i")) for _ in range(arr.ndim))
    shape = tuple(arr.shape)
    ctx = EffectBlob(stmt=Query(Table(tns, idxs), Table(Literal(arr), idxs)))
    return LazyTensor(tns, ctx, shape, arr.fill, arr.element_type, arr.device)

einop(prgm, /, **kwargs)

Execute an einsum expression using the specified array framework.

This function parses and executes einsum-like expressions with extended syntax that supports various operations beyond traditional Einstein summation notation.

Parameters:

Name	Type	Description	Default
prgm	str	Einsum program string specifying the computation. The syntax supports: - Assignment: "C[i,j] = A[i,j] + B[j,i]" - Increment: "C[i,j] += A[i,k] * B[k,j]" - Reductions: "C[i] += A[i,j]", "C[i] max= A[i,j]", "C[i] &= A[i,j]" - Arithmetic operations: +, -, *, /, //, %, ** - Comparison operations: ==, !=, <, <=, >, >= - Logical operations: and, or, not - Bitwise operations: &, |, ^, <<, >> - Function calls and complex expressions with parentheses - Mathematical functions: abs, sqrt, exp, log, sin, cos, tan, etc. - Literal values: integers, floats, booleans, and complex numbers - Python operator precedence and parentheses for grouping	required
**kwargs	Any	Named arrays referenced in the einsum expression. The keys should match the tensor names used in the program string.	{}

Returns:

Type	Description
Any	The result array from executing the einsum expression.

Examples:

>>> rng = np.random.default_rng(42)
>>> A = rng.integers(0, 10, size=(3, 4))
>>> B = rng.integers(0, 10, size=(4, 3))
>>> # Matrix addition with transpose
>>> C = einop("C[i,j] = A[i,j] + B[j,i]", A=A, B=B)
>>> # Matrix multiplication
>>> D = einop("D[i,j] += A[i,k] * B[k,j]", A=A, B=B)
>>> # Min-Plus multiplication with shift
>>> E = einop("E[i] min= A[k,i] + D[k,j] << 1", A=A, D=D)

Source code in src/finch/interface/eager.py

def einop(prgm: str, /, **kwargs: Any) -> Any:
    """Execute an einsum expression using the specified array framework.

    This function parses and executes einsum-like expressions with extended syntax
    that supports various operations beyond traditional Einstein summation notation.

    Args:
        prgm (str): Einsum program string specifying the computation. The syntax
            supports:
            - Assignment: "C[i,j] = A[i,j] + B[j,i]"
            - Increment: "C[i,j] += A[i,k] * B[k,j]"
            - Reductions: "C[i] += A[i,j]", "C[i] max= A[i,j]", "C[i] &= A[i,j]"
            - Arithmetic operations: +, -, *, /, //, %, **
            - Comparison operations: ==, !=, <, <=, >, >=
            - Logical operations: and, or, not
            - Bitwise operations: &, |, ^, <<, >>
            - Function calls and complex expressions with parentheses
            - Mathematical functions: abs, sqrt, exp, log, sin, cos, tan, etc.
            - Literal values: integers, floats, booleans, and complex numbers
            - Python operator precedence and parentheses for grouping
        **kwargs: Named arrays referenced in the einsum expression. The keys
            should match the tensor names used in the program string.

    Returns:
        The result array from executing the einsum expression.

    Examples:
        >>> rng = np.random.default_rng(42)
        >>> A = rng.integers(0, 10, size=(3, 4))
        >>> B = rng.integers(0, 10, size=(4, 3))
        >>> # Matrix addition with transpose
        >>> C = einop("C[i,j] = A[i,j] + B[j,i]", A=A, B=B)
        >>> # Matrix multiplication
        >>> D = einop("D[i,j] += A[i,k] * B[k,j]", A=A, B=B)
        >>> # Min-Plus multiplication with shift
        >>> E = einop("E[i] min= A[k,i] + D[k,j] << 1", A=A, D=D)
    """
    if builtins.any(isinstance(v, lazy.LazyTensor) for v in kwargs.values()):
        return lazy.einop(prgm, **kwargs)
    return compute(lazy.einop(prgm, **kwargs))

einsum(*args, **kwargs)

einsum(subscripts, *operands)

Evaluates the Einstein summation convention on the operands.

Using the Einstein summation convention, many common multi-dimensional, linear algebraic array operations can be represented in a simple fashion. In implicit mode einsum computes these values.

In explicit mode, einsum provides further flexibility to compute other array operations that might not be considered classical Einstein summation operations, by disabling, or forcing summation over specified subscript labels.

See the notes and examples for clarification.

Parameters

subscripts : str Specifies the subscripts for summation as comma separated list of subscript labels. An implicit (classical Einstein summation) calculation is performed unless the explicit indicator '->' is included as well as subscript labels of the precise output form. operands : list of array_like These are the arrays for the operation.

Returns

output : ndarray The calculation based on the Einstein summation convention.

Notes

The Einstein summation convention can be used to compute many multi-dimensional, linear algebraic array operations. einsum provides a succinct way of representing these.

A non-exhaustive list of these operations, which can be computed by einsum, is shown below along with examples:

• Trace of an array, :py:func:numpy.trace.
• Return a diagonal, :py:func:numpy.diag.
• Array axis summations, :py:func:numpy.sum.
• Transpositions and permutations, :py:func:numpy.transpose.
• Matrix multiplication and dot product, :py:func:numpy.matmul :py:func:numpy.dot.
• Vector inner and outer products, :py:func:numpy.inner :py:func:numpy.outer.
• Broadcasting, element-wise and scalar multiplication, :py:func:numpy.multiply.
• Tensor contractions, :py:func:numpy.tensordot.
• Chained array operations, in efficient calculation order, :py:func:numpy.einsum_path.

The subscripts string is a comma-separated list of subscript labels, where each label refers to a dimension of the corresponding operand. Whenever a label is repeated it is summed, so np.einsum('i,i', a, b) is equivalent to :py:func:np.inner(a,b) <numpy.inner>. If a label appears only once, it is not summed, so np.einsum('i', a) produces a view of a with no changes. A further example np.einsum('ij,jk', a, b) describes traditional matrix multiplication and is equivalent to :py:func:np.matmul(a,b) <numpy.matmul>. Repeated subscript labels in one operand take the diagonal. For example, np.einsum('ii', a) is equivalent to :py:func:np.trace(a) <numpy.trace>.

In implicit mode, the chosen subscripts are important since the axes of the output are reordered alphabetically. This means that np.einsum('ij', a) doesn't affect a 2D array, while np.einsum('ji', a) takes its transpose. Additionally, np.einsum('ij,jk', a, b) returns a matrix multiplication, while, np.einsum('ij,jh', a, b) returns the transpose of the multiplication since subscript 'h' precedes subscript 'i'.

In explicit mode the output can be directly controlled by specifying output subscript labels. This requires the identifier '->' as well as the list of output subscript labels. This feature increases the flexibility of the function since summing can be disabled or forced when required. The call np.einsum('i->', a) is like :py:func:np.sum(a) <numpy.sum> if a is a 1-D array, and np.einsum('ii->i', a) is like :py:func:np.diag(a) <numpy.diag> if a is a square 2-D array. The difference is that einsum does not allow broadcasting by default. Additionally np.einsum('ij,jh->ih', a, b) directly specifies the order of the output subscript labels and therefore returns matrix multiplication, unlike the example above in implicit mode.

To enable and control broadcasting, use an ellipsis. Default NumPy-style broadcasting is done by adding an ellipsis to the left of each term, like np.einsum('...ii->...i', a). np.einsum('...i->...', a) is like :py:func:np.sum(a, axis=-1) <numpy.sum> for array a of any shape. To take the trace along the first and last axes, you can do np.einsum('i...i', a), or to do a matrix-matrix product with the left-most indices instead of rightmost, one can do np.einsum('ij...,jk...->ik...', a, b).

einsum also provides an alternative way to provide the subscripts and operands as einsum(op0, sublist0, op1, sublist1, ..., [sublistout]). If the output shape is not provided in this format einsum will be calculated in implicit mode, otherwise it will be performed explicitly. The examples below have corresponding einsum calls with the two parameter methods.

Examples

a = np.arange(25).reshape(5, 5) b = np.arange(5) c = np.arange(6).reshape(2, 3)

Trace of a matrix:

np.einsum("ii", a) np.int64(60) np.einsum(a, [0, 0]) np.int64(60) np.trace(a) np.int64(60)

Extract the diagonal (requires explicit form):

np.einsum("ii->i", a) array([ 0, 6, 12, 18, 24]) np.einsum(a, [0, 0], [0]) array([ 0, 6, 12, 18, 24]) np.diag(a) array([ 0, 6, 12, 18, 24])

Sum over an axis (requires explicit form):

np.einsum("ij->i", a) array([ 10, 35, 60, 85, 110]) np.einsum(a, [0, 1], [0]) array([ 10, 35, 60, 85, 110]) np.sum(a, axis=1) array([ 10, 35, 60, 85, 110])

For higher dimensional arrays summing a single axis can be done with ellipsis:

np.einsum("...j->...", a) array([ 10, 35, 60, 85, 110]) np.einsum(a, [Ellipsis, 1], [Ellipsis]) array([ 10, 35, 60, 85, 110])

Compute a matrix transpose, or reorder any number of axes:

np.einsum("ji", c) array([[0, 3], [1, 4], [2, 5]]) np.einsum("ij->ji", c) array([[0, 3], [1, 4], [2, 5]]) np.einsum(c, [1, 0]) array([[0, 3], [1, 4], [2, 5]]) np.transpose(c) array([[0, 3], [1, 4], [2, 5]])

Vector inner products:

np.einsum("i,i", b, b) np.int64(30) np.einsum(b, [0], b, [0]) np.int64(30) np.inner(b, b) np.int64(30)

Matrix vector multiplication:

np.einsum("ij,j", a, b) array([ 30, 80, 130, 180, 230]) np.einsum(a, [0, 1], b, [1]) array([ 30, 80, 130, 180, 230]) np.dot(a, b) array([ 30, 80, 130, 180, 230]) np.einsum("...j,j", a, b) array([ 30, 80, 130, 180, 230])

Broadcasting and scalar multiplication:

np.einsum("..., ...", 3, c) array([[ 0, 3, 6], [ 9, 12, 15]]) np.einsum(",ij", 3, c) array([[ 0, 3, 6], [ 9, 12, 15]]) np.einsum(3, [Ellipsis], c, [Ellipsis]) array([[ 0, 3, 6], [ 9, 12, 15]]) np.multiply(3, c) array([[ 0, 3, 6], [ 9, 12, 15]])

Vector outer product:

np.einsum("i,j", np.arange(2) + 1, b) array([[0, 1, 2, 3, 4], [0, 2, 4, 6, 8]]) np.einsum(np.arange(2) + 1, [0], b, [1]) array([[0, 1, 2, 3, 4], [0, 2, 4, 6, 8]]) np.outer(np.arange(2) + 1, b) array([[0, 1, 2, 3, 4], [0, 2, 4, 6, 8]])

Tensor contraction:

a = np.arange(60.0).reshape(3, 4, 5) b = np.arange(24.0).reshape(4, 3, 2) np.einsum("ijk,jil->kl", a, b) array([[4400., 4730.], [4532., 4874.], [4664., 5018.], [4796., 5162.], [4928., 5306.]]) np.einsum(a, [0, 1, 2], b, [1, 0, 3], [2, 3]) array([[4400., 4730.], [4532., 4874.], [4664., 5018.], [4796., 5162.], [4928., 5306.]]) np.tensordot(a, b, axes=([1, 0], [0, 1])) array([[4400., 4730.], [4532., 4874.], [4664., 5018.], [4796., 5162.], [4928., 5306.]])

Example of ellipsis use:

a = np.arange(6).reshape((3, 2)) b = np.arange(12).reshape((4, 3)) np.einsum("ki,jk->ij", a, b) array([[10, 28, 46, 64], [13, 40, 67, 94]]) np.einsum("ki,...k->i...", a, b) array([[10, 28, 46, 64], [13, 40, 67, 94]]) np.einsum("k...,jk", a, b) array([[10, 28, 46, 64], [13, 40, 67, 94]])

Source code in src/finch/interface/eager.py

def einsum(*args, **kwargs):
    """
    einsum(subscripts, *operands)

    Evaluates the Einstein summation convention on the operands.

    Using the Einstein summation convention, many common multi-dimensional,
    linear algebraic array operations can be represented in a simple fashion.
    In *implicit* mode `einsum` computes these values.

    In *explicit* mode, `einsum` provides further flexibility to compute
    other array operations that might not be considered classical Einstein
    summation operations, by disabling, or forcing summation over specified
    subscript labels.

    See the notes and examples for clarification.

    Parameters
    ----------
    subscripts : str
        Specifies the subscripts for summation as comma separated list of
        subscript labels. An implicit (classical Einstein summation)
        calculation is performed unless the explicit indicator '->' is
        included as well as subscript labels of the precise output form.
    operands : list of array_like
        These are the arrays for the operation.

    Returns
    -------
    output : ndarray
        The calculation based on the Einstein summation convention.

    Notes
    -----
    The Einstein summation convention can be used to compute
    many multi-dimensional, linear algebraic array operations. `einsum`
    provides a succinct way of representing these.

    A non-exhaustive list of these operations,
    which can be computed by `einsum`, is shown below along with examples:

    * Trace of an array, :py:func:`numpy.trace`.
    * Return a diagonal, :py:func:`numpy.diag`.
    * Array axis summations, :py:func:`numpy.sum`.
    * Transpositions and permutations, :py:func:`numpy.transpose`.
    * Matrix multiplication and dot product, :py:func:`numpy.matmul`
        :py:func:`numpy.dot`.
    * Vector inner and outer products, :py:func:`numpy.inner`
        :py:func:`numpy.outer`.
    * Broadcasting, element-wise and scalar multiplication,
        :py:func:`numpy.multiply`.
    * Tensor contractions, :py:func:`numpy.tensordot`.
    * Chained array operations, in efficient calculation order,
        :py:func:`numpy.einsum_path`.

    The subscripts string is a comma-separated list of subscript labels,
    where each label refers to a dimension of the corresponding operand.
    Whenever a label is repeated it is summed, so ``np.einsum('i,i', a, b)``
    is equivalent to :py:func:`np.inner(a,b) <numpy.inner>`. If a label
    appears only once, it is not summed, so ``np.einsum('i', a)``
    produces a view of ``a`` with no changes. A further example
    ``np.einsum('ij,jk', a, b)`` describes traditional matrix multiplication
    and is equivalent to :py:func:`np.matmul(a,b) <numpy.matmul>`.
    Repeated subscript labels in one operand take the diagonal.
    For example, ``np.einsum('ii', a)`` is equivalent to
    :py:func:`np.trace(a) <numpy.trace>`.

    In *implicit mode*, the chosen subscripts are important
    since the axes of the output are reordered alphabetically.  This
    means that ``np.einsum('ij', a)`` doesn't affect a 2D array, while
    ``np.einsum('ji', a)`` takes its transpose. Additionally,
    ``np.einsum('ij,jk', a, b)`` returns a matrix multiplication, while,
    ``np.einsum('ij,jh', a, b)`` returns the transpose of the
    multiplication since subscript 'h' precedes subscript 'i'.

    In *explicit mode* the output can be directly controlled by
    specifying output subscript labels.  This requires the
    identifier '->' as well as the list of output subscript labels.
    This feature increases the flexibility of the function since
    summing can be disabled or forced when required. The call
    ``np.einsum('i->', a)`` is like :py:func:`np.sum(a) <numpy.sum>`
    if ``a`` is a 1-D array, and ``np.einsum('ii->i', a)``
    is like :py:func:`np.diag(a) <numpy.diag>` if ``a`` is a square 2-D array.
    The difference is that `einsum` does not allow broadcasting by default.
    Additionally ``np.einsum('ij,jh->ih', a, b)`` directly specifies the
    order of the output subscript labels and therefore returns matrix
    multiplication, unlike the example above in implicit mode.

    To enable and control broadcasting, use an ellipsis.  Default
    NumPy-style broadcasting is done by adding an ellipsis
    to the left of each term, like ``np.einsum('...ii->...i', a)``.
    ``np.einsum('...i->...', a)`` is like
    :py:func:`np.sum(a, axis=-1) <numpy.sum>` for array ``a`` of any shape.
    To take the trace along the first and last axes,
    you can do ``np.einsum('i...i', a)``, or to do a matrix-matrix
    product with the left-most indices instead of rightmost, one can do
    ``np.einsum('ij...,jk...->ik...', a, b)``.

    `einsum` also provides an alternative way to provide the subscripts and
    operands as ``einsum(op0, sublist0, op1, sublist1, ..., [sublistout])``.
    If the output shape is not provided in this format `einsum` will be
    calculated in implicit mode, otherwise it will be performed explicitly.
    The examples below have corresponding `einsum` calls with the two
    parameter methods.

    Examples
    --------
    >>> a = np.arange(25).reshape(5, 5)
    >>> b = np.arange(5)
    >>> c = np.arange(6).reshape(2, 3)

    Trace of a matrix:

    >>> np.einsum("ii", a)
    np.int64(60)
    >>> np.einsum(a, [0, 0])
    np.int64(60)
    >>> np.trace(a)
    np.int64(60)

    Extract the diagonal (requires explicit form):

    >>> np.einsum("ii->i", a)
    array([ 0,  6, 12, 18, 24])
    >>> np.einsum(a, [0, 0], [0])
    array([ 0,  6, 12, 18, 24])
    >>> np.diag(a)
    array([ 0,  6, 12, 18, 24])

    Sum over an axis (requires explicit form):

    >>> np.einsum("ij->i", a)
    array([ 10,  35,  60,  85, 110])
    >>> np.einsum(a, [0, 1], [0])
    array([ 10,  35,  60,  85, 110])
    >>> np.sum(a, axis=1)
    array([ 10,  35,  60,  85, 110])

    For higher dimensional arrays summing a single axis can be done
    with ellipsis:

    >>> np.einsum("...j->...", a)
    array([ 10,  35,  60,  85, 110])
    >>> np.einsum(a, [Ellipsis, 1], [Ellipsis])
    array([ 10,  35,  60,  85, 110])

    Compute a matrix transpose, or reorder any number of axes:

    >>> np.einsum("ji", c)
    array([[0, 3],
           [1, 4],
           [2, 5]])
    >>> np.einsum("ij->ji", c)
    array([[0, 3],
           [1, 4],
           [2, 5]])
    >>> np.einsum(c, [1, 0])
    array([[0, 3],
           [1, 4],
           [2, 5]])
    >>> np.transpose(c)
    array([[0, 3],
           [1, 4],
           [2, 5]])

    Vector inner products:

    >>> np.einsum("i,i", b, b)
    np.int64(30)
    >>> np.einsum(b, [0], b, [0])
    np.int64(30)
    >>> np.inner(b, b)
    np.int64(30)

    Matrix vector multiplication:

    >>> np.einsum("ij,j", a, b)
    array([ 30,  80, 130, 180, 230])
    >>> np.einsum(a, [0, 1], b, [1])
    array([ 30,  80, 130, 180, 230])
    >>> np.dot(a, b)
    array([ 30,  80, 130, 180, 230])
    >>> np.einsum("...j,j", a, b)
    array([ 30,  80, 130, 180, 230])

    Broadcasting and scalar multiplication:

    >>> np.einsum("..., ...", 3, c)
    array([[ 0,  3,  6],
           [ 9, 12, 15]])
    >>> np.einsum(",ij", 3, c)
    array([[ 0,  3,  6],
           [ 9, 12, 15]])
    >>> np.einsum(3, [Ellipsis], c, [Ellipsis])
    array([[ 0,  3,  6],
           [ 9, 12, 15]])
    >>> np.multiply(3, c)
    array([[ 0,  3,  6],
           [ 9, 12, 15]])

    Vector outer product:

    >>> np.einsum("i,j", np.arange(2) + 1, b)
    array([[0, 1, 2, 3, 4],
           [0, 2, 4, 6, 8]])
    >>> np.einsum(np.arange(2) + 1, [0], b, [1])
    array([[0, 1, 2, 3, 4],
           [0, 2, 4, 6, 8]])
    >>> np.outer(np.arange(2) + 1, b)
    array([[0, 1, 2, 3, 4],
           [0, 2, 4, 6, 8]])

    Tensor contraction:

    >>> a = np.arange(60.0).reshape(3, 4, 5)
    >>> b = np.arange(24.0).reshape(4, 3, 2)
    >>> np.einsum("ijk,jil->kl", a, b)
    array([[4400., 4730.],
           [4532., 4874.],
           [4664., 5018.],
           [4796., 5162.],
           [4928., 5306.]])
    >>> np.einsum(a, [0, 1, 2], b, [1, 0, 3], [2, 3])
    array([[4400., 4730.],
           [4532., 4874.],
           [4664., 5018.],
           [4796., 5162.],
           [4928., 5306.]])
    >>> np.tensordot(a, b, axes=([1, 0], [0, 1]))
    array([[4400., 4730.],
           [4532., 4874.],
           [4664., 5018.],
           [4796., 5162.],
           [4928., 5306.]])

    Example of ellipsis use:

    >>> a = np.arange(6).reshape((3, 2))
    >>> b = np.arange(12).reshape((4, 3))
    >>> np.einsum("ki,jk->ij", a, b)
    array([[10, 28, 46, 64],
           [13, 40, 67, 94]])
    >>> np.einsum("ki,...k->i...", a, b)
    array([[10, 28, 46, 64],
           [13, 40, 67, 94]])
    >>> np.einsum("k...,jk", a, b)
    array([[10, 28, 46, 64],
           [13, 40, 67, 94]])
    """

    if builtins.any(isinstance(v, lazy.LazyTensor) for v in args):
        return lazy.einsum(*args, **kwargs)
    return compute(lazy.einsum(*args, **kwargs))

element(fill_type=None, element_type=None, position_type=None, buffer_factory=NumpyBufferFType, buffer_type=None)

Creates an ElementLevelFType with the given parameters.

Parameters:

Name	Type	Description	Default
fill_type	Any	The fill of the level's format, or the value of a static fill.	None
element_type	FType | None	The type of elements stored in the level.	None
position_type	FType | None	The type of positions within the level.	None
buffer_factory	Any	The factory used to create buffers for the level.	NumpyBufferFType
buffer_type	Any	Format of the value stored in the level.	None

Returns:

Type	Description
ElementLevelFType	An instance of ElementLevelFType.

Source code in src/finch/tensor/level/element_level.py

def element(
    fill_type: Any = None,
    element_type: FType | None = None,
    position_type: FType | None = None,
    buffer_factory: Any = NumpyBufferFType,
    buffer_type: Any = None,
) -> ElementLevelFType:
    """
    Creates an ElementLevelFType with the given parameters.

    Args:
        fill_type: The fill of the level's format, or the value of a static
            fill.
        element_type: The type of elements stored in the level.
        position_type: The type of positions within the level.
        buffer_factory: The factory used to create buffers for the level.
        buffer_type: Format of the value stored in the level.

    Returns:
        An instance of ElementLevelFType.
    """
    match fill_type:
        case AbstractFillFType() | None:
            pass
        case AbstractFill():
            raise TypeError(f"element takes the fill of a format, not {fill_type}")
        case _:
            fill_type = StaticFillFType(fill_type)
    return ElementLevelFType(
        fill_type=fill_type,
        element_type=element_type,
        position_type=position_type,
        buffer_factory=buffer_factory,
        buffer_type=buffer_type,
    )

fisinstance(x, f)

Check if x is an instance of f.

Source code in src/finch/algebra/ftypes.py

def fisinstance(x, f: FType):
    """
    Check if `x` is an instance of `f`.
    """
    return f.fisinstance(x)

from_binsparse(array)

Convert a Binsparse tensor from the reference binsparse package into a Finch tensor which shares its buffers wherever possible.

Dense levels become one Finch DenseLevel per dimension, rank-1 sparse levels become SparseListLevel, and higher-rank sparse levels become SparseCOOLevel.

Source code in src/finch/fileio/binsparse.py

def from_binsparse(array: bsp.BinsparseTensor) -> FiberTensor:
    """
    Convert a Binsparse tensor from the reference `binsparse` package into a
    Finch tensor which shares its buffers wherever possible.

    Dense levels become one Finch `DenseLevel` per dimension, rank-1 sparse
    levels become `SparseListLevel`, and higher-rank sparse levels become
    `SparseCOOLevel`.
    """
    if not isinstance(array, bsp.CustomTensor):
        container = bsp.InMemoryBinsparseContainer()
        array.serialize(container, alias=False)
        array = bsp.CustomTensor.parse(container)
    match array:
        case bsp.CustomTensor(
            shape=shape, transpose=None, level=bsp.BinsparseLevel() as level
        ):
            values = _element_values(level)
            fill_value = array.fill_value if array.fill is True else 0
            element_format = element(
                fill_type=StaticFillFType(values.dtype.type(fill_value)),
                element_type=ftype(values.dtype),
                position_type=ftype(_position_dtype(level)),
            )
            return FiberTensor(_from_level(level, shape, element_format, 0))
        case _:
            raise NotImplementedError(
                "Finch does not support transposed Binsparse tensors."
            )

ftype(x)

ftype(x: FinchOperator) -> FinchOperatorFType

ftype(x: Tensor) -> TensorFType

ftype(x: tuple) -> TupleFType

ftype(x: NumpyBuffer) -> NumpyBufferFType

ftype(x: np.integer) -> FDTypeInteger

ftype(x: FT) -> FT

ftype(x: FTyped[FT]) -> FT

ftype(x: Any) -> FType

Return the corresponding FType for a given object. Recognizes numpy, Python builtins, and Python tuples. Sometimes recognizes types. Override .ftype to customize the ftype of an object.

Source code in src/finch/algebra/ftypes.py

def ftype(x: Any) -> FType:
    """Return the corresponding FType for a given object.  Recognizes numpy,
    Python builtins, and Python tuples.  Sometimes recognizes types.
    Override .ftype to customize the ftype of an object.
    """
    if isinstance(x, FType):
        return x
    if isinstance(x, FTyped):
        return x.ftype
    dtype: Any = None
    if isinstance(x, np.dtype):
        if x.fields is not None:
            assert x.names is not None
            return TupleFType.from_tuple(
                tuple(ftype(x.fields[name][0]) for name in x.names)
            )
        dtype = x
    elif isinstance(x, np.void) and x.dtype.fields is not None:
        assert x.dtype.names is not None
        return TupleFType.from_tuple(
            tuple(ftype(x.dtype.fields[name][0]) for name in x.dtype.names)
        )
    elif isinstance(x, np.generic):
        dtype = x.dtype
    elif isinstance(x, type) and issubclass(x, np.generic):
        dtype = np.dtype(x)
    if dtype is not None:
        if dtype == np.dtype(np.bool_):
            return bool
        if dtype == np.dtype(np.int8):
            return int8
        if dtype == np.dtype(np.int16):
            return int16
        if dtype == np.dtype(np.int32):
            return int32
        if dtype == np.dtype(np.int64):
            return int64
        if dtype == np.dtype(np.uint8):
            return uint8
        if dtype == np.dtype(np.uint16):
            return uint16
        if dtype == np.dtype(np.uint32):
            return uint32
        if dtype == np.dtype(np.uint64):
            return uint64
        if dtype == np.dtype(np.float16):
            return float16
        if dtype == np.dtype(np.float32):
            return float32
        if dtype == np.dtype(np.float64):
            return float64
        if dtype == np.dtype(np.complex64):
            return complex64
        if dtype == np.dtype(np.complex128):
            return complex128
        raise NotImplementedError
    if type(x) is builtins.bool or x is builtins.bool:
        return bool_
    if type(x) is builtins.int or x is builtins.int:
        return int_
    if type(x) is builtins.float or x is builtins.float:
        return float_
    if type(x) is builtins.complex or x is builtins.complex:
        return complex_
    if type(x) is builtins.str or x is builtins.str:
        return str_
    if x is None:
        return none_
    if isinstance(x, tuple):
        T = type(x)
        if hasattr(T, "_fields") and all(isinstance(field, str) for field in T._fields):
            return NamedTupleFType(
                T.__name__,
                [(fieldname, ftype(getattr(x, fieldname))) for fieldname in T._fields],
            )
        return TupleFType.from_tuple(tuple(ftype(elem) for elem in x))
    raise NotImplementedError

fused(f, /, ctx=None)

• fused(f): A decorator that marks a function as fused. This allows the function to be used with the fuse function for automatic fusion of operations.

Parameters: - f: The function to be marked as fused.

Returns: - A wrapper function that applies the fusion mechanism to the original function.

Source code in src/finch/interface/fuse.py

def fused(f, /, ctx=None):
    """
    - fused(f):
    A decorator that marks a function as fused. This allows the function to be used with
    the `fuse` function for automatic fusion of operations.

    Parameters:
    - f: The function to be marked as fused.

    Returns:
    - A wrapper function that applies the fusion mechanism to the original function.
    """
    if ctx is None:
        ctx = get_default_scheduler()

    def wrapper(*args):
        return fuse(f, *args, ctx=ctx)

    return wrapper

isdtype(dtype, kind)

Check if a dtype belongs to an Array API dtype kind.

Source code in src/finch/algebra/ftypes.py

def isdtype(dtype, kind):
    """
    Check if a dtype belongs to an Array API dtype kind.
    """
    dtype = ftype(dtype)
    if isinstance(kind, tuple):
        return any(isdtype(dtype, k) for k in kind)
    if isinstance(kind, str):
        match kind:
            case "bool":
                return builtins.bool(dtype == bool)
            case "signed integer":
                return isinstance(dtype, FDTypeSignedInteger)
            case "unsigned integer":
                return isinstance(dtype, FDTypeUnsignedInteger)
            case "integral":
                return isinstance(dtype, FDTypeInteger)
            case "real floating":
                return isinstance(dtype, FDTypeFloat) and not isinstance(
                    dtype, FDTypeComplex
                )
            case "complex floating":
                return isinstance(dtype, FDTypeComplex)
            case "numeric":
                return isinstance(dtype, FDTypeNumeric)
            case _:
                raise ValueError(f"Unrecognized dtype kind: {kind!r}")
    return builtins.bool(dtype == ftype(kind))

jit(f, /, ctx=None)

A decorator that marks a function for just-in-time compilation. This allows the function to be compiled and optimized for performance when called.

• f: The function to be marked for JIT compilation. This function can use basic python control flow and operations (e.g. while, for, if). However, it shouldn't use more complex features like generators, classes, or recursion.
• ctx: The scheduler to use for computation. Defaults to the result of get_default_scheduler().

• A transformed function that inserts defer and compute statements to do tracing and optimization.

Example usage: @jit def my_function(A, B, C): D = A @ B while some_condition(D): D = D + C return D

In this example, my_function will be transformed to include defer and compute statements, allowing it to be optimized and executed efficiently when called. def opt_my_function(A, B, C): A, B = defer(A), defer(B) D = A @ B D = compute(D) while some_condition(D): C = defer(C) D = D + C D = compute(D) D = compute(D) return D

Source code in src/finch/finch_fused/jit.py

def jit(f, /, ctx=None):
    """
    A decorator that marks a function for just-in-time compilation. This allows the
    function to be compiled and optimized for performance when called.

    Parameters:
    - f: The function to be marked for JIT compilation. This function can use
        basic python control flow and operations (e.g. while, for, if). However,
        it shouldn't use more complex features like generators, classes, or recursion.
    - ctx: The scheduler to use for computation. Defaults to the result of
        `get_default_scheduler()`.

    Returns:
    - A transformed function that inserts defer and compute statements to do tracing
       and optimization.


    Example usage:
    @jit
    def my_function(A, B, C):
        D = A @ B
        while some_condition(D):
            D = D + C
        return D

    In this example, `my_function` will be transformed to include defer and compute
    statements, allowing it to be optimized and executed efficiently when called.
    def opt_my_function(A, B, C):
        A, B = defer(A), defer(B)
        D = A @ B
        D = compute(D)
        while some_condition(D):
            C = defer(C)
            D = D + C
            D = compute(D)
        D = compute(D)
        return D
    """
    if ctx is None:
        ctx = get_default_scheduler()
    fused_fn = parse_fused_function(f)
    transformed_fn = insert_lazy_and_compute(fused_fn)
    return fused_function_to_python_function(transformed_fn)

matmul(x1, x2)

Computes the matrix product.

Returns a LazyTensor if either x1 or x2 is a LazyTensor. Otherwise, computes the result eagerly.

Source code in src/finch/interface/eager.py

def matmul(x1, x2, /):
    """
    Computes the matrix product.

    Returns a LazyTensor if either x1 or x2 is a LazyTensor.
    Otherwise, computes the result eagerly.
    """
    if isinstance(x1, lazy.LazyTensor) or isinstance(x2, lazy.LazyTensor):
        return lazy.matmul(x1, x2)
    c = lazy.matmul(x1, x2)
    return compute(c)

matrix_transpose(x)

Computes the transpose of a matrix or stack of matrices.

Source code in src/finch/interface/eager.py

def matrix_transpose(x, /):
    """
    Computes the transpose of a matrix or stack of matrices.
    """
    if isinstance(x, lazy.LazyTensor):
        return lazy.matrix_transpose(x)
    return compute(lazy.matrix_transpose(x))

moveaxis(x, source, destination)

Moves array axes (dimensions) to new positions, while leaving other axes in their original positions.

Args

• x (array) - input array.
• source - Axes to move.
• destination - indices defining the desired positions for each respective source axis index.

Returns

• out (array) - an array containing reordered axes.

Source code in src/finch/interface/eager.py

def moveaxis(x, source: int | tuple[int, ...], destination: int | tuple[int, ...], /):
    """
    Moves array axes (dimensions) to new positions,
    while leaving other axes in their original positions.

    Args
    ---------
    - x (array) - input array.
    - source - Axes to move.
    - destination - indices defining the desired
    positions for each respective source axis index.

    Returns
    --------
    - out (array) - an array containing reordered axes.
    """
    if isinstance(x, lazy.LazyTensor):
        return lazy.moveaxis(x, source, destination)
    return compute(lazy.moveaxis(x, source, destination))

tensordot(x1, x2, /, *, axes=2)

Computes the tensordot operation.

Returns a LazyTensor if either x1 or x2 is a LazyTensor. Otherwise, computes the result eagerly.

Source code in src/finch/interface/eager.py

def tensordot(x1, x2, /, *, axes: int | tuple[Sequence[int], Sequence[int]] = 2):
    """
    Computes the tensordot operation.

    Returns a LazyTensor if either x1 or x2 is a LazyTensor.
    Otherwise, computes the result eagerly.
    """
    if isinstance(x1, lazy.LazyTensor) or isinstance(x2, lazy.LazyTensor):
        return lazy.tensordot(x1, x2, axes=axes)
    return compute(lazy.tensordot(x1, x2, axes=axes))

to_binsparse(tns)

Convert a Finch tensor to a Binsparse tensor which shares its buffers.

Source code in src/finch/fileio/binsparse.py

def to_binsparse(tns: FiberTensor) -> bsp.CustomTensor:
    """Convert a Finch tensor to a Binsparse tensor which shares its buffers."""
    level, shape = _to_level(tns.lvl, 0)
    return bsp.CustomTensor(
        shape=shape,
        number_of_stored_values=_element_values(level).size,
        fill=True,
        fill_value=tns.fill_value,
        level=level,
    )

vecdot(x1, x2, /, *, axis=-1)

Computes the (vector) dot product of two arrays.

Parameters

x1: array The first input tensor. x2: array The second input tensor. axis: int, optional The axis along which to compute the dot product. Default is -1 (last axis).

Returns

out: array A tensor containing the dot product of x1 and x2 along the specified axis.

Source code in src/finch/interface/eager.py

def vecdot(x1, x2, /, *, axis=-1):
    """
    Computes the (vector) dot product of two arrays.

    Parameters
    ----------
    x1: array
        The first input tensor.
    x2: array
        The second input tensor.
    axis: int, optional
        The axis along which to compute the dot product. Default is -1 (last axis).

    Returns
    -------
    out: array
        A tensor containing the dot product of `x1` and `x2` along the specified axis.
    """
    if isinstance(x1, lazy.LazyTensor) or isinstance(x2, lazy.LazyTensor):
        return lazy.vecdot(x1, x2, axis=axis)
    return compute(lazy.vecdot(x1, x2, axis=axis))