fiber_tensor

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)

fiber_tensor(lvl)

Creates a FiberTensorFType with the given level ftype and position type.

Parameters:

Name	Type	Description	Default
lvl	LevelFType	The level ftype to be used for the tensor.	required

Returns: An instance of a fiber tensor format.

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

def fiber_tensor(lvl: LevelFType) -> FiberTensorFType:
    """
    Creates a FiberTensorFType with the given level ftype and position type.

    Args:
        lvl: The level ftype to be used for the tensor.
    Returns:
        An instance of a fiber tensor format.
    """
    return FiberTensorFType(lvl)