bufferized_ndarray

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)

BufferizedNDArrayAccessor

Bases: Tensor

A class representing a tensor view that is bufferized. This is used to create a view of a tensor with a specific extent.

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

class BufferizedNDArrayAccessor(Tensor):
    """
    A class representing a tensor view that is bufferized.
    This is used to create a view of a tensor with a specific extent.
    """

    def __init__(self, tns: BufferizedNDArray, nind=None, pos=None, op=None):
        self.tns = tns
        if pos is None:
            tns_ftype = cast(BufferizedNDArrayFType, ftype(self.tns))
            pos = tns_ftype.buf_t.length_type(0)
        self.pos = pos
        self.op = op
        if nind is None:
            nind = 0
        self.nind = nind

    @property
    def lvl(self):
        return BufferizedNDArrayLevel(self.tns, self.nind)

    @property
    def ftype(self):
        return FiberTensorFType(self.lvl.ftype, self.tns.device)

    @property
    def shape(self):
        return self.tns.shape[self.nind :]

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

    @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 access(self, indices, op):
        if len(indices) + self.nind > self.tns.ndim:
            raise IndexError(
                f"Too many indices for tensor access: "
                f"got {len(indices)} indices for tensor with "
                f"{self.tns.ndim - self.nind} dimensions."
            )
        for i, idx in enumerate(indices):
            if not (0 <= idx < self.tns.shape[self.nind + i]):
                raise IndexError(
                    f"Index {idx} out of bounds for axis {self.nind + i} "
                    f"with size {self.tns.shape[self.nind + i]}"
                )
        pos = self.pos
        for i, idx in enumerate(indices):
            pos += idx * self.tns.strides[self.nind + i]
        return BufferizedNDArrayAccessor(self.tns, self.nind + len(indices), pos, op)

    def unwrap(self):
        """
        Unwrap the tensor view to get the underlying tensor.
        This is used to get the original tensor from a tensor view.
        """
        assert self.ndim == 0, "Cannot unwrap a tensor view with non-zero dimension."
        return self.tns.val.load(self.pos)

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

    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.")

    def increment(self, val):
        """
        Increment the tensor view with a value.
        This updates the tensor at the specified index with the operation and value.
        """
        if self.op is None:
            raise ValueError("No operation defined for increment.")
        assert self.ndim == 0, "Cannot unwrap a tensor view with non-zero dimension."
        if isinstance(val, Scalar):
            val = val.item()
        self.tns.val.store(self.pos, self.op(self.tns.val.load(self.pos), val))
        return self

element_type property

Data type of the tensor elements.

fill_value property

Default value to fill the tensor.

shape_type property

Shape type of the tensor.

increment(val)

Increment the tensor view with a value. This updates the tensor at the specified index with the operation and value.

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

def increment(self, val):
    """
    Increment the tensor view with a value.
    This updates the tensor at the specified index with the operation and value.
    """
    if self.op is None:
        raise ValueError("No operation defined for increment.")
    assert self.ndim == 0, "Cannot unwrap a tensor view with non-zero dimension."
    if isinstance(val, Scalar):
        val = val.item()
    self.tns.val.store(self.pos, self.op(self.tns.val.load(self.pos), val))
    return self

unwrap()

Unwrap the tensor view to get the underlying tensor. This is used to get the original tensor from a tensor view.

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

def unwrap(self):
    """
    Unwrap the tensor view to get the underlying tensor.
    This is used to get the original tensor from a tensor view.
    """
    assert self.ndim == 0, "Cannot unwrap a tensor view with non-zero dimension."
    return self.tns.val.load(self.pos)

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__}"
    )

BufferizedNDArrayLevel dataclass

Bases: Level

A level cursor sharing an ndarray's storage and shape metadata.

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

@dataclass
class BufferizedNDArrayLevel(Level):
    """A level cursor sharing an ndarray's storage and shape metadata."""

    tns: BufferizedNDArray
    nind: int = 0

    def with_fill(self, fill_value: AbstractFill):
        return replace(self, tns=self.tns.with_fill(fill_value))

    @property
    def ftype(self):
        return BufferizedNDArrayLevelFType(self.tns.ftype, self.nind)

    @property
    def shape(self):
        return self.tns.shape[self.nind :]

    @property
    def stride(self):
        return self.tns.strides[self.nind] if self.ndim else self.position_type(1)

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

    @property
    def lvl(self):
        if self.ndim == 0:
            raise TypeError("Scalar ndarray levels have no child level")
        return BufferizedNDArrayLevel(self.tns, self.nind + 1)