Reduction Operations

simplegrad.functions.reduction.sum(x: Tensor, dim: int | None = None) -> Tensor

Sum tensor elements along a dimension.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required
dim	int | None	Dimension to reduce. If None, sums all elements.	None

Returns:

Type	Description
Tensor	Reduced tensor (keepdims=True).

Source code in simplegrad/functions/reduction.py

def sum(x: Tensor, dim: int | None = None) -> Tensor:
    """Sum tensor elements along a dimension.

    Args:
        x: Input tensor.
        dim: Dimension to reduce. If None, sums all elements.

    Returns:
        Reduced tensor (keepdims=True).
    """
    return _Sum.apply(x, dim, oper=f"sum(d={dim})")

simplegrad.functions.reduction.mean(x: Tensor, dim: int | None = None) -> Tensor

Compute the mean of tensor elements along a dimension.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required
dim	int | None	Dimension to reduce. If None, averages all elements.	None

Returns:

Type	Description
Tensor	Reduced tensor.

Source code in simplegrad/functions/reduction.py

def mean(x: Tensor, dim: int | None = None) -> Tensor:
    """Compute the mean of tensor elements along a dimension.

    Args:
        x: Input tensor.
        dim: Dimension to reduce. If None, averages all elements.

    Returns:
        Reduced tensor.
    """
    if dim is None:
        n = 1
        for s in x.shape:
            n *= s
        return sum(x) / n
    return sum(x, dim=dim) / x.shape[dim]

simplegrad.functions.reduction.trace(x: Tensor) -> Tensor

Compute the trace (sum of diagonal elements) of a square matrix.

Parameters:

Name	Type	Description	Default
x	Tensor	2D square tensor.	required

Returns:

Type	Description
Tensor	Scalar tensor containing the trace.

Raises:

Type	Description
ValueError	If x is not a 2D square tensor (checked in eager mode).

Source code in simplegrad/functions/reduction.py

def trace(x: Tensor) -> Tensor:
    """Compute the trace (sum of diagonal elements) of a square matrix.

    Args:
        x: 2D square tensor.

    Returns:
        Scalar tensor containing the trace.

    Raises:
        ValueError: If x is not a 2D square tensor (checked in eager mode).
    """
    if x.values is not None and (x.values.ndim != 2 or x.values.shape[0] != x.values.shape[1]):
        raise ValueError("Trace is only defined for square matrices")
    return _Trace.apply(x)

simplegrad.functions.reduction.argmax(x: Tensor, dim: int | None = None, dtype: str = 'int32') -> Tensor

Return indices of maximum values along a dimension.

Not differentiable — comp_grad is always False on the output.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required
dim	int | None	Dimension to reduce. If None, returns the flat index.	None
dtype	str	Integer dtype for the output indices.	'int32'

Returns:

Type	Description
Tensor	Integer tensor of argmax indices.

Source code in simplegrad/functions/reduction.py

def argmax(x: Tensor, dim: int | None = None, dtype: str = "int32") -> Tensor:
    """Return indices of maximum values along a dimension.

    Not differentiable — ``comp_grad`` is always False on the output.

    Args:
        x: Input tensor.
        dim: Dimension to reduce. If None, returns the flat index.
        dtype: Integer dtype for the output indices.

    Returns:
        Integer tensor of argmax indices.
    """
    out = _Argmax.apply(x, dim, dtype, oper=f"argmax(d={dim})")
    out.dtype = get_dtype_class(dtype)
    return out

simplegrad.functions.reduction.argmin(x: Tensor, dim: int | None = None, dtype: str = 'int32') -> Tensor

Return indices of minimum values along a dimension.

Not differentiable — comp_grad is always False on the output.

Parameters:

Name	Type	Description	Default
x	Tensor	Input tensor.	required
dim	int | None	Dimension to reduce. If None, returns the flat index.	None
dtype	str	Integer dtype for the output indices.	'int32'

Returns:

Type	Description
Tensor	Integer tensor of argmin indices.

Source code in simplegrad/functions/reduction.py

def argmin(x: Tensor, dim: int | None = None, dtype: str = "int32") -> Tensor:
    """Return indices of minimum values along a dimension.

    Not differentiable — ``comp_grad`` is always False on the output.

    Args:
        x: Input tensor.
        dim: Dimension to reduce. If None, returns the flat index.
        dtype: Integer dtype for the output indices.

    Returns:
        Integer tensor of argmin indices.
    """
    out = _Argmin.apply(x, dim, dtype, oper=f"argmin(d={dim})")
    out.dtype = get_dtype_class(dtype)
    return out