lazy

EffectBlob

The EffectBlob class represents a collection of prior queries in a lazy evaluation context, which can be extracted into a linear trace of LogicStatements (in the order they were issued). Note that the global effect_stamp counter is required to maintain the relative ordering of effect blobs in the face of in-place mutations to tensors. Each effect blob contains a logic statement and references to other effect blobs. The datastructure is designed to be immutable, mergeable in O(1) time, and support linear time extraction. To achieve linear time, we use a hash table to avoid re-extracting already seen blobs.

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

class EffectBlob:
    """
    The EffectBlob class represents a collection of prior queries in a lazy
    evaluation context, which can be extracted into a linear trace of
    LogicStatements (in the order they were issued). Note that the global
    `effect_stamp` counter is required to maintain the relative ordering of
    effect blobs in the face of in-place mutations to tensors.  Each effect blob
    contains a logic statement and references to other effect blobs. The
    datastructure is designed to be immutable, mergeable in O(1) time, and
    support linear time extraction. To achieve linear time, we use a hash table
    to avoid re-extracting already seen blobs.
    """

    stamp: int
    stmt: LogicStatement
    blobs: tuple[EffectBlob, ...]

    def __init__(
        self,
        stmt: LogicStatement | None = None,
        blobs: tuple[EffectBlob, ...] | None = None,
    ):
        global effect_stamp
        if stmt is None:
            stmt = Plan()
        if blobs is None:
            blobs = ()
        self.stmt = stmt
        self.blobs = blobs

        try:
            curr_stamp = effect_stamp.value
        except AttributeError:
            effect_stamp.value = 0
            curr_stamp = 0

        self.stamp = curr_stamp
        effect_stamp.value += 1

    def exec(self, stmt: LogicStatement) -> EffectBlob:
        return EffectBlob(stmt=stmt, blobs=(self,))

    def eval(self, ex: LogicExpression) -> tuple[Alias, EffectBlob]:
        var = Alias(gensym("A"))
        return var, self.exec(Query(Table(var, ex.fields()), ex))

    def join(self, *blobs: EffectBlob) -> EffectBlob:
        return EffectBlob(blobs=(self, *blobs))

    def trace(self) -> tuple[LogicStatement, ...]:
        stmts = list[tuple[int, LogicStatement]]()
        seen = set[int]()
        self._trace(seen, stmts)
        stmts.sort()
        return tuple(stmt for _, stmt in stmts)

    def _trace(self, seen: set[int], stmts: list[tuple[int, LogicStatement]]) -> None:
        if id(self) not in seen:
            seen.add(id(self))
            stmts.append((self.stamp, self.stmt))
            for blob in self.blobs:
                blob._trace(seen, stmts)

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: Any,
        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 = as_fill(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_value=self._fill_value,
            _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."
    )

LazyTensorFType

Bases: TensorFType

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

class LazyTensorFType(TensorFType):
    _fill_value: AbstractFill
    _element_type: FType
    _shape_type: tuple[FType, ...]

    def __init__(
        self,
        _fill_value: Any,
        _element_type: FType | type,
        _shape_type: tuple[FType | type, ...],
        _device=None,
    ):
        self._fill_value = as_fill(_fill_value)
        self._element_type = ftype(_element_type)
        self._shape_type = tuple(ftype(dim_t) for dim_t in _shape_type)
        self._device = normalize_device(_device)

    def __eq__(self, other):
        if not isinstance(other, LazyTensorFType):
            return False
        return (
            builtins.bool(np.all(ffuncs.same(self._fill_value, other._fill_value)))
            and self._element_type == other._element_type
            and self._shape_type == other._shape_type
            and self.device == other.device
        )

    def __hash__(self):
        return hash(
            (
                ffuncs.samehash(self._fill_value),
                self._element_type,
                self._shape_type,
                self.device,
            )
        )

    def construct(self, shape: tuple) -> LazyTensor:
        idxs = tuple(Field(gensym("i")) for _ in shape)
        ctx = EffectBlob()
        expr = Table(Literal(FillTensor(shape, self._fill_value)), idxs)
        data, ctx = ctx.eval(expr)
        return LazyTensor(
            data=data,
            ctx=ctx,
            shape=shape,
            fill_value=self._fill_value,
            element_type=self._element_type,
            device=self.device,
        )

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

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

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

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

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

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

    def from_numpy(self, arr):
        raise NotImplementedError

__call__(val)

Convert a tensor to this tensor type.

Parameters:

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

Returns: A LazyTensor instance of this type.

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

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

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

all(x, /, *, axis=None, keepdims=False, init=None)

Test whether all elements of input array x along given axis are True.

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

def all(
    x,
    /,
    *,
    axis: int | tuple[int, ...] | None = None,
    keepdims: bool = False,
    init=None,
):
    """
    Test whether all elements of input array ``x`` along given axis are True.
    """
    x = defer(x)
    return reduce(
        ffuncs.and_,
        elementwise(ffuncs.truth, x),
        axis=axis,
        keepdims=keepdims,
        init=init,
    )

any(x, /, *, axis=None, keepdims=False, init=None)

Test whether any element of input array x along given axis is True.

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

def any(
    x,
    /,
    *,
    axis: int | tuple[int, ...] | None = None,
    keepdims: bool = False,
    init=None,
):
    """
    Test whether any element of input array ``x`` along given axis is True.
    """
    x = defer(x)
    return reduce(
        ffuncs.or_,
        elementwise(ffuncs.truth, x),
        axis=axis,
        keepdims=keepdims,
        init=init,
    )

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(*arrays)

Broadcasts one or more arrays against one another.

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

def broadcast_arrays(*arrays: LazyTensor) -> tuple[LazyTensor, ...]:
    """
    Broadcasts one or more arrays against one another.
    """
    shape = _broadcast_shape(*(array.shape for array in arrays))
    return tuple(broadcast_to(arr, shape) for arr in arrays)

broadcast_to(tensor, /, shape)

Broadcasts a lazy tensor to a specified shape.

Parameters:

Name	Type	Description	Default
tensor	Any	The lazy tensor to broadcast.	required
shape	tuple	The target shape to broadcast to.	required

Returns:

Type	Description
LazyTensor	A new lazy tensor with the specified shape.

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

def broadcast_to(tensor: Any, /, shape: tuple) -> LazyTensor:
    """
    Broadcasts a lazy tensor to a specified shape.

    Args:
        tensor: The lazy tensor to broadcast.
        shape: The target shape to broadcast to.

    Returns:
        A new lazy tensor with the specified shape.
    """
    tensor = defer(tensor)

    if not is_broadcastable_directional(tensor.shape, shape):
        raise ValueError(
            f"Tensor with shape {tensor.shape} is not broadcastable "
            f"to the shape {shape}"
        )
    return elementwise(ffuncs.first_arg, tensor, FillTensor(shape, np.False_))

conj(x)

Computes the complex conjugate of the input tensor x.

Parameters

x: LazyTensor The input tensor to compute the complex conjugate of.

Returns

LazyTensor A new LazyTensor with the complex conjugate of x.

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

def conj(x) -> LazyTensor:
    """
    Computes the complex conjugate of the input tensor `x`.

    Parameters
    ----------
    x: LazyTensor
        The input tensor to compute the complex conjugate of.

    Returns
    -------
    LazyTensor
        A new LazyTensor with the complex conjugate of `x`.
    """
    return elementwise(ffuncs.conjugate, x)

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.ftype.fill_value, arr.element_type, arr.device
    )

elementwise(f, *args)

elementwise(f, *args) -> LazyTensor:

Applies the function f elementwise to the given arguments, following broacasting rules. The function f should be a callable that takes the same number of arguments as the number of tensors passed to elementwise.

The function will automatically handle broadcasting of the input tensors to ensure they have compatible shapes. For example, elementwise(ffunc.add, x, y) is equivalent to x + y.

Parameters: - f: The function to apply elementwise. - *args: The tensors to apply the function to. These tensors should be compatible for broadcasting.

Returns: - LazyTensor: The tensor, out, of results from applying f elementwise to the input tensors. After broadcasting the arguments to the same shape, for each index i, out[*i] = f(args[0][*i], args[1][*i], ...).

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

def elementwise(f: FinchOperator, *args) -> LazyTensor:
    """
        elementwise(f, *args) -> LazyTensor:

    Applies the function f elementwise to the given arguments, following
    [broacasting rules](https://numpy.org/doc/stable/user/basics.broadcasting.html).
    The function f should be a callable that takes the same number of arguments
    as the number of tensors passed to `elementwise`.

    The function will automatically handle broadcasting of the input tensors to
    ensure they have compatible shapes.  For example, `elementwise(ffunc.add,
    x, y)` is equivalent to `x + y`.

    Parameters:
    - f: The function to apply elementwise.
    - *args: The tensors to apply the function to. These tensors should be
        compatible for broadcasting.

    Returns:
    - LazyTensor: The tensor, `out`, of results from applying `f` elementwise to
    the input tensors.  After broadcasting the arguments to the same shape, for
    each index `i`, `out[*i] = f(args[0][*i], args[1][*i], ...)`.
    """
    args = tuple(defer(a) for a in args)
    shapes = tuple(a.shape for a in args)
    shape = _broadcast_shape(*shapes)
    ndim = len(shape)
    idxs = tuple(Field(gensym("i")) for _ in range(ndim))
    bargs: list[LogicExpression] = []
    for arg, arg_shape in zip(args, shapes, strict=True):
        arg_ndim = len(arg_shape)
        idims = []
        odims = []
        for i in range(ndim - arg_ndim, ndim):
            if arg_shape[i - ndim + arg_ndim] == shape[i]:
                idims.append(idxs[i])
                odims.append(idxs[i])
            else:
                if arg_shape[i - ndim + arg_ndim] != 1:
                    raise ValueError("Invalid shape for broadcasting")
                idims.append(Field(gensym("j")))
        bargs.append(Reorder(Table(arg.data, tuple(idims)), tuple(odims)))
    expr = Reorder(MapJoin(Literal(f), tuple(bargs)), idxs)
    new_fill_value = apply_fill(f, *[a.ftype.fill_value for a in args])
    op_type = f.ftype
    if not isinstance(op_type, CallableFType):
        raise TypeError(f"Expected a callable type, got {op_type}")
    new_element_type = return_type(op_type, *[a.element_type for a in args])
    ctx = args[0].ctx.join(*[x.ctx for x in args[1:]])
    data, ctx = ctx.eval(expr)
    return LazyTensor(
        data,
        ctx,
        shape,
        new_fill_value,
        new_element_type,
        common_device(*(t.device for t in args)),
    )

expand_dims(x, /, axis=0)

Expands the shape of an array by inserting a new axis (dimension) of size one at the position specified by axis.

Parameters

x: array input array. axis: int axis position (zero-based). If x has rank (i.e, number of dimensions) N, a valid axis must reside on the closed-interval [-N-1, N]. If provided a negative axis, the axis position at which to insert a singleton dimension must be computed as N + axis + 1. Hence, if provided -1, the resolved axis position must be N (i.e., a singleton dimension must be appended to the input array x). If provided -N-1, the resolved axis position must be 0 (i.e., a singleton dimension must be prepended to the input array x).

Returns

out: array an expanded output array having the same data type as x.

Raises

IndexError If provided an invalid axis position, an IndexError should be raised.

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

def expand_dims(
    x,
    /,
    axis: int | tuple[int, ...] = 0,
) -> LazyTensor:
    """
    Expands the shape of an array by inserting a new axis (dimension) of size one at the
    position specified by ``axis``.

    Parameters
    ----------
    x: array
        input array.
    axis: int
        axis position (zero-based). If ``x`` has rank (i.e, number of dimensions) ``N``,
        a valid ``axis`` must reside on the closed-interval ``[-N-1, N]``. If provided a
        negative ``axis``, the axis position at which to insert a singleton dimension
        must be computed as ``N + axis + 1``. Hence, if provided ``-1``, the resolved
        axis position must be ``N`` (i.e., a singleton dimension must be appended to the
        input array ``x``). If provided ``-N-1``, the resolved axis position must be
        ``0`` (i.e., a singleton dimension must be prepended to the input array ``x``).

    Returns
    -------
    out: array
        an expanded output array having the same data type as ``x``.

    Raises
    ------
    IndexError
        If provided an invalid ``axis`` position, an ``IndexError`` should be raised.
    """
    x = defer(x)
    if not isinstance(axis, (tuple, list)):
        axis = (axis,)
    axis = normalize_axis_tuple(axis, x.ndim + len(axis))
    if isinstance(axis, int):
        raise IndexError(
            f"Invalid axis: {axis}. Axis must be an integer or a tuple of integers."
        )
    if not len(axis) == len(set(axis)):
        raise IndexError("axis must be unique")
    if not set(axis).issubset(range(x.ndim + len(axis))):
        raise IndexError(
            f"Invalid axis: {axis}. Axis must be unique and must be in the range "
            f"[-{x.ndim + len(axis) - 1}, {x.ndim + len(axis) - 1}]."
        )
    offset = [0] * (x.ndim + len(axis))
    for d in axis:
        offset[d] = 1
    offset = list(accumulate(offset))
    idxs_1 = tuple(Field(gensym("i")) for _ in range(x.ndim))
    idxs_2 = tuple(
        Field(gensym("i")) if n in axis else idxs_1[n - offset[n]]
        for n in range(x.ndim + len(axis))
    )
    expr = Reorder(Table(x.data, idxs_1), idxs_2)
    shape_2 = tuple(
        1 if n in axis else x.shape[n - offset[n]] for n in range(x.ndim + len(axis))
    )
    data_2, ctx = x.ctx.eval(expr)
    return LazyTensor(data_2, ctx, shape_2, x.fill_value, x.element_type, x.device)

full(shape, fill_value, *, dtype=None, device=None)

Returns a new array having a specified shape and filled with fill_value.

Parameters: - shape (Union[int, Tuple[int, ...]]): output array shape. - fill_value (Union[bool, int, float, complex]): fill value. - dtype (Optional[dtype]): output array data type. If dtype is None, the output array data type must be inferred from fill_value according to the following rules: * If the fill value is an int, the output array data type must be the default integer data type. * If the fill value is a float, the output array data type must be the default real-valued floating-point data type. * If the fill value is a complex number, the output array data type must be the default complex floating-point data type. * If the fill value is a bool, the output array must have a boolean data type. Default: None.

Returns:

• out (array): an array where every element is equal to fill_value.

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

def full(
    shape: int | tuple[int, ...],
    fill_value: bool | complex,
    *,
    dtype: Any | None = None,
    device=None,
):
    """
    Returns a new array having a specified shape and filled with fill_value.

    Parameters:
    - shape (Union[int, Tuple[int, ...]]): output array shape.
    - fill_value (Union[bool, int, float, complex]): fill value.
    - dtype (Optional[dtype]): output array data type. If dtype is None, the
    output array data type must be inferred from fill_value according to the
    following rules:
        * If the fill value is an int, the output array data type must be the
            default integer data type.
        * If the fill value is a float, the output array data type must be the
            default real-valued floating-point data type.
        * If the fill value is a complex number, the output array data type must
            be the default complex floating-point data type.
        * If the fill value is a bool, the output array must have a boolean data
            type. Default: None.

    Returns:

    - out (array): an array where every element is equal to fill_value.
    """
    shape = (shape,) if isinstance(shape, int) else tuple(shape)
    return broadcast_to(asarray(fill_value, dtype=dtype, device=device), shape)

is_broadcastable(shape_a, shape_b)

Returns True if shape_a and shape_b are broadcastable according to numpy rules.

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

def is_broadcastable(shape_a, shape_b):
    """
    Returns True if shape_a and shape_b are broadcastable according to numpy rules.
    """
    for a, b in zip_longest(reversed(shape_a), reversed(shape_b), fillvalue=1):
        if a != b and a != 1 and b != 1:
            return False
    return True

is_broadcastable_directional(shape_a, shape_b)

Returns True if shape_a is broadcastable to shape_b according to numpy rules. This is a directional check, so it allows only shape_a to be changed

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

def is_broadcastable_directional(shape_a, shape_b):
    """
    Returns True if shape_a is broadcastable to shape_b according to numpy rules.
    This is a directional check, so it allows only shape_a to be changed
    """
    for a, b in zip_longest(reversed(shape_a), reversed(shape_b), fillvalue=1):
        if a != b and a != 1:
            return False
    return True

matmul(x1, x2)

Performs matrix multiplication between two tensors.

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

def matmul(x1, x2) -> LazyTensor:
    """
    Performs matrix multiplication between two tensors.
    """

    def _matmul_helper(a, b) -> LazyTensor:
        """
        For arrays greater than 1D
        """
        if a.ndim < 2 or b.ndim < 2:
            raise ValueError("Both inputs must be at least 2D arrays")
        if a.shape[-1] != b.shape[-2]:
            raise ValueError("Dimensions mismatch for matrix multiplication")
        # check all preceeding dimensions match
        batch_a, batch_b = a.shape[:-2], b.shape[:-2]
        if not is_broadcastable(batch_a, batch_b):
            raise ValueError(
                "Batch dimensions are not broadcastable for matrix multiplication"
            )
        return reduce(
            ffuncs.add,
            multiply(expand_dims(a, axis=-1), expand_dims(b, axis=-3)),
            axis=-2,
        )

    x1 = defer(x1)
    x2 = defer(x2)

    if x1.ndim < 1 or x2.ndim < 1:
        raise ValueError("Both inputs must be at least 1D arrays")

    if x1.ndim == 1 and x2.ndim == 1:
        return reduce(ffuncs.add, multiply(x1, x2), axis=0)

    if x1.ndim == 1:
        x1 = expand_dims(x1, axis=0)  # make it a row vector
        result = _matmul_helper(x1, x2)
        return squeeze(result, axis=-2)  # remove the prepended singleton dimension

    if x2.ndim == 1:
        x2 = expand_dims(x2, axis=1)  # make it a column vector
        result = _matmul_helper(x1, x2)
        return squeeze(result, axis=-1)  # remove the appended singleton dimension

    return _matmul_helper(x1, x2)

matrix_transpose(x)

Transposes the input tensor x.

Parameters

x: LazyTensor The input tensor to be transposed. Must have at least 2 dimensions.

Returns

LazyTensor A new LazyTensor with the axes of x transposed.

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

def matrix_transpose(x) -> LazyTensor:
    """
    Transposes the input tensor `x`.

    Parameters
    ----------
    x: LazyTensor
        The input tensor to be transposed. Must have at least 2 dimensions.

    Returns
    -------
    LazyTensor
        A new LazyTensor with the axes of `x` transposed.
    """
    x = defer(x)
    if x.ndim < 2:
        # this is following numpy's behavior.
        # data-apis specification assumes that input is atleast 2D
        raise ValueError(
            "Input tensor must have at least 2 dimensions for transposition"
        )
    # swap the last two axes
    return permute_dims(x, axes=(*range(x.ndim - 2), x.ndim - 1, x.ndim - 2))

max(x, /, *, axis=None, keepdims=False, init=None)

Return the maximum of input array arr along given axis.

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

def max(
    x,
    /,
    *,
    axis: int | tuple[int, ...] | None = None,
    keepdims: bool = False,
    init=None,
):
    """
    Return the maximum of input array ``arr`` along given axis.
    """
    x = defer(x)
    return reduce(ffuncs.max, x, axis=axis, keepdims=keepdims, init=init)

mean(x, /, *, axis=None, keepdims=False)

Calculates the arithmetic mean of the input array x.

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

def mean(x, /, *, axis: int | tuple[int, ...] | None = None, keepdims: bool = False):
    """
    Calculates the arithmetic mean of the input array ``x``.
    """
    x = defer(x)
    n = (
        np.prod(tuple(x.shape[i] for i in range(x.ndim) if i in axis))
        if isinstance(axis, tuple)
        else (np.prod(x.shape) if axis is None else x.shape[axis])
    )
    s = sum(x, axis=axis, keepdims=keepdims)
    return truediv(s, x.element_type(n))

min(x, /, *, axis=None, keepdims=False, init=None)

Return the minimum of input array arr along given axis.

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

def min(
    x,
    /,
    *,
    axis: int | tuple[int, ...] | None = None,
    keepdims: bool = False,
    init=None,
):
    """
    Return the minimum of input array ``arr`` along given axis.
    """
    x = defer(x)
    return reduce(ffuncs.min, x, axis=axis, keepdims=keepdims, init=init)

moveaxis(x, source, destination)

Moves axes of an array to new positions.

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

def moveaxis(x, source: int | tuple[int, ...], destination: int | tuple[int, ...], /):
    """
    Moves axes of an array to new positions.
    """
    x = defer(x)

    # argument validation
    # handles uniqueness, int -> tuple, and bound check
    source = normalize_axis_tuple(source, x.ndim, "source")
    destination = normalize_axis_tuple(destination, x.ndim, "destination")

    if len(source) != len(destination):
        raise ValueError("Source and Destination indices must have the same length")

    final_order = [i for i in range(x.ndim) if i not in source]
    for dest, src in sorted(zip(destination, source, strict=True)):
        final_order.insert(dest, src)

    return permute_dims(x, axes=tuple(final_order))

permute_dims(arg, /, axes)

Permutes the axes (dimensions) of an array x.

Parameters

x: array input array. axes: Tuple[int, ...] tuple containing a permutation of (0, 1, ..., N-1) where N is the number of axes (dimensions) of x.

Returns

out: array an array containing the axes permutation. The returned array must have the same data type as x.

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

def permute_dims(arg, /, axes: tuple[int, ...]) -> LazyTensor:
    """
    Permutes the axes (dimensions) of an array ``x``.

    Parameters
    ----------
    x: array
        input array.
    axes: Tuple[int, ...]
        tuple containing a permutation of ``(0, 1, ..., N-1)`` where ``N`` is the number
        of axes (dimensions) of ``x``.

    Returns
    -------
    out: array
        an array containing the axes permutation. The returned array must have the same
        data type as ``x``.
    """
    arg = defer(arg)
    axis = normalize_axis_tuple(axes, arg.ndim + len(axes))
    idxs = tuple(Field(gensym("i")) for _ in range(arg.ndim))
    expr = Reorder(Table(arg.data, idxs), tuple(idxs[i] for i in axis))
    data, ctx = arg.ctx.eval(expr)
    return LazyTensor(
        data,
        ctx,
        tuple(arg.shape[i] for i in axis),
        arg.fill_value,
        arg.element_type,
        device=arg.device,
    )

reduce(op, x, /, *, axis=None, dtype=None, keepdims=False, init=None)

Reduces the input array x with the binary operator op. Reduces along the specified axis, with an initial value init.

Parameters

x: array input array. Should have a numeric data type. axis: Optional[Union[int, Tuple[int, ...]]] axis or axes along which reduction must be computed. By default, the reduction must be computed over the entire array. If a tuple of integers, reductions must be computed over multiple axes. Default: None.

Optional[dtype]

data type of the returned array. If None, a suitable data type will be calculated.

bool

if True, the reduced axes (dimensions) must be included in the result as singleton dimensions, and, accordingly, the result must be compatible with the input array (see :ref:broadcasting). Otherwise, if False, the reduced axes (dimensions) must not be included in the result. Default: False.

Optional

Initial value for the reduction. If None, a suitable initial value will be calculated. The initial value must be compatible with the operation defined by op. For example, if op is addition, the initial value should be zero; if op is multiplication, the initial value should be one.

Returns

out: array If the reduction was computed over the entire array, a zero-dimensional array containing the reduction; otherwise, a non-zero-dimensional array containing the reduction. The returned array must have a data type as described by the dtype parameter above.

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

def reduce(
    op: FinchOperator,
    x,
    /,
    *,
    axis: int | tuple[int, ...] | None = None,
    dtype=None,
    keepdims: bool = False,
    init=None,
):
    """
    Reduces the input array ``x`` with the binary operator ``op``. Reduces along
    the specified `axis`, with an initial value `init`.

    Parameters
    ----------
    x: array
        input array. Should have a numeric data type.
    axis: Optional[Union[int, Tuple[int, ...]]]
        axis or axes along which reduction must be computed. By default, the reduction
        must be computed over the entire array. If a tuple of integers, reductions must
        be computed over multiple axes. Default: ``None``.

    dtype: Optional[dtype]
        data type of the returned array. If ``None``, a suitable data type will be
        calculated.

    keepdims: bool
        if ``True``, the reduced axes (dimensions) must be included in the result as
        singleton dimensions, and, accordingly, the result must be compatible with the
        input array (see :ref:`broadcasting`). Otherwise, if ``False``, the reduced axes
        (dimensions) must not be included in the result. Default: ``False``.

    init: Optional
        Initial value for the reduction. If ``None``, a suitable initial value will be
        calculated. The initial value must be compatible with the operation defined by
        ``op``. For example, if ``op`` is addition, the initial value should be zero; if
        ``op`` is multiplication, the initial value should be one.

    Returns
    -------
    out: array
        If the reduction was computed over the entire array, a zero-dimensional array
        containing the reduction; otherwise, a non-zero-dimensional array containing the
        reduction. The returned array must have a data type as described by the
        ``dtype`` parameter above.
    """
    x = defer(x)
    op_type = op.ftype
    if not isinstance(op_type, FinchOperatorFType):
        raise TypeError(f"Expected a Finch operator type, got {op_type}")
    explicit_dtype = dtype is not None
    if dtype is not None:
        dtype = ftype(dtype)
    if init is None:
        init = init_value(op_type, dtype or x.element_type)
        if explicit_dtype:
            assert dtype is not None
            init = dtype(init)
    if axis is None:
        axis = tuple(range(x.ndim))
    axis = normalize_axis_tuple(axis, x.ndim)
    if axis is None or isinstance(axis, int):
        raise ValueError("axis must be a tuple")

    shape = tuple(x.shape[n] for n in range(x.ndim) if n not in axis)
    fields = tuple(Field(gensym("i")) for _ in range(x.ndim))
    data: LogicExpression = Aggregate(
        Literal(op),
        Literal(init),
        Table(x.data, fields),
        tuple(fields[i] for i in axis),
    )
    if keepdims:
        keeps = tuple(
            fields[i] if i not in axis else Field(gensym("j")) for i in range(x.ndim)
        )
        data = Reorder(data, keeps)
        shape = tuple(x.shape[i] if i not in axis else 1 for i in range(x.ndim))
    if dtype is None:
        dtype = fixpoint_type(op_type, init, x.element_type)
    expr, ctx = x.ctx.eval(data)
    return LazyTensor(expr, ctx, shape, init, dtype, x.device)

squeeze(x, /, axis)

Removes singleton dimensions (axes) from x.

Parameters

x: array input array. axis: Union[int, Tuple[int, ...]] axis (or axes) to squeeze.

Returns

out: array an output array having the same data type and elements as x.

Raises

ValueError If a specified axis has a size greater than one (i.e., it is not a singleton dimension), a ValueError should be raised.

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

def squeeze(
    x,
    /,
    axis: int | tuple[int, ...],
) -> LazyTensor:
    """
    Removes singleton dimensions (axes) from ``x``.

    Parameters
    ----------
    x: array
        input array.
    axis: Union[int, Tuple[int, ...]]
        axis (or axes) to squeeze.

    Returns
    -------
    out: array
        an output array having the same data type and elements as ``x``.

    Raises
    ------
    ValueError
        If a specified axis has a size greater than one (i.e., it is not a
        singleton dimension), a ``ValueError`` should be raised.
    """
    x = defer(x)
    if not isinstance(axis, (tuple, list)):
        axis = (axis,)
    axis = normalize_axis_tuple(axis, x.ndim)
    if isinstance(axis, int):
        raise ValueError(f"Invalid axis: {axis}. Axis must be a tuple of integers.")
    if len(axis) != len(set(axis)):
        raise ValueError(f"Invalid axis: {axis}. Axis must be unique.")
    if not set(axis).issubset(range(x.ndim)):
        raise ValueError(f"Invalid axis: {axis}. Axis must be within bounds.")
    if not builtins.all(x.shape[d] == 1 for d in axis):
        raise ValueError(f"Invalid axis: {axis}. Axis to drop must have size 1.")
    newaxis = [n for n in range(x.ndim) if n not in axis]
    idxs_1 = tuple(Field(gensym("i")) for _ in range(x.ndim))
    idxs_2 = tuple(idxs_1[n] for n in newaxis)
    expr = Reorder(Table(x.data, idxs_1), idxs_2)
    shape_2 = tuple(x.shape[n] for n in newaxis)
    data_2, ctx = x.ctx.eval(expr)
    return LazyTensor(data_2, ctx, shape_2, x.fill_value, x.element_type, x.device)

std(x, /, *, axis=None, correction=0.0, keepdims=False)

Calculates the standard deviation of the input array x.

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

def std(
    x,
    /,
    *,
    axis: int | tuple[int, ...] | None = None,
    correction: float = 0.0,
    keepdims: bool = False,
):
    """
    Calculates the standard deviation of the input array ``x``.
    """
    x = defer(x)
    n = (
        np.prod(tuple(x.shape[i] for i in range(x.ndim) if i in axis))
        if isinstance(axis, tuple)
        else (np.prod(x.shape) if axis is None else x.shape[axis])
    )
    m = mean(x, axis=axis, keepdims=True)
    norm = vector_norm(x - m, axis=axis, keepdims=keepdims, ord=2)
    return truediv(norm, sqrt(norm.element_type(n - correction)))

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/lazy.py

def tensordot(
    x1, x2, /, *, axes: int | tuple[Sequence[int], Sequence[int]] = 2
) -> LazyTensor:
    """
    Computes the tensordot operation.

    Returns a LazyTensor if either x1 or x2 is a LazyTensor.
    Otherwise, computes the result eagerly.
    """
    x1 = defer(x1)
    x2 = defer(x2)

    # Parse axes
    if not isinstance(axes, tuple):
        N = int(axes)
        if N < 0:
            raise ValueError("expected axes to be a non-negative integer")
        axes_a = list(range(x1.ndim - N, x1.ndim))
        axes_b = list(range(N))
    else:
        axes_a, axes_b = (list(ax) for ax in axes)

    # Normalize negative axes. We need list
    axes_a = [(a if a >= 0 else x1.ndim + a) for a in axes_a]
    axes_b = [(b if b >= 0 else x2.ndim + b) for b in axes_b]

    # Check axes lengths and shapes
    if len(axes_a) != len(axes_b):
        raise ValueError("shape-mismatch for sum")
    for a, b in zip(axes_a, axes_b, strict=True):
        if x1.shape[a] != x2.shape[b]:
            raise ValueError("shape-mismatch for sum")

    # Move axes to contract to the end of x1 and to the front of x2
    notin_a = [k for k in range(x1.ndim) if k not in axes_a]
    notin_b = [k for k in range(x2.ndim) if k not in axes_b]
    newaxes_a = notin_a + axes_a
    newaxes_b = axes_b + notin_b
    # Permute
    x1p = permute_dims(x1, tuple(newaxes_a))
    x2p = permute_dims(x2, tuple(newaxes_b))

    # Expand x1p and x2p so that their contracted axes align for broadcasting
    # so we can multiply them

    # For x1p, add len(notin_b) singleton dims at the end
    added_dims = tuple(-(i + 1) for i in range(len(notin_b)))
    x1p = expand_dims(x1p, axis=added_dims)

    # For x2p, add len(notin_a) singleton dims at the front
    added_dims = tuple(i for i in range(len(notin_a)))
    x2p = expand_dims(x2p, axis=added_dims)

    # Multiply (broadcasted)
    expanded_product = multiply(x1p, x2p)

    sum_axes = tuple(range(len(notin_a), len(notin_a) + len(axes_a)))
    if not sum_axes:
        return expanded_product
    return sum(
        expanded_product,
        axis=sum_axes,
        dtype=result_type(x1.element_type, x2.element_type),
    )

var(x, /, *, axis=None, correction=0.0, keepdims=False)

Calculates the variance of the input array x.

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

def var(
    x,
    /,
    *,
    axis: int | tuple[int, ...] | None = None,
    correction: float = 0.0,
    keepdims: bool = False,
):
    """
    Calculates the variance of the input array ``x``.
    """
    x = defer(x)
    n = (
        np.prod(tuple(x.shape[i] for i in range(x.ndim) if i in axis))
        if isinstance(axis, tuple)
        else (np.prod(x.shape) if axis is None else x.shape[axis])
    )
    m = mean(x, axis=axis, keepdims=True)
    v = sum(pow(x - m, x.element_type(2.0)), axis=axis, keepdims=keepdims)
    return truediv(v, x.element_type(n - correction))

vecdot(x1, x2, /, *, axis=-1)

Computes the vector dot product along the specified axis.

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

def vecdot(x1, x2, /, *, axis=-1) -> LazyTensor:
    """
    Computes the vector dot product along the specified axis.
    """
    x1 = defer(x1)
    x2 = defer(x2)

    # check broadcastability
    if not is_broadcastable(x1.shape, x2.shape):
        raise ValueError("Shapes are not broadcastable for vector dot product")

    # check if dims of axis are the same
    if x1.shape[axis] != x2.shape[axis]:
        raise ValueError(
            "Shapes are not compatible for vector dot product along the specified axis"
        )

    return reduce(
        ffuncs.add,
        multiply(conj(x1), x2),
        axis=axis,
        dtype=result_type(x1.element_type, x2.element_type),
    )