Function normalize_reduction_dim

backends/mlx/ops.py:326–355

Normalize dim argument for reduction operations. Extracts and normalizes the dim argument from handler args, returning a list of axes and the keepdim flag. Handles both list-based dims (e.g., sum.dim_IntList) and single int dims (e.g., prod.dim_int). Args: args: The ha

(
    args: List[Any], start_idx: int = 1
)

def normalize_reduction_dim(
    args: List[Any], start_idx: int = 1
) -> Tuple[Optional[List[int]], bool]:
    """
    Normalize dim argument for reduction operations.

    Extracts and normalizes the dim argument from handler args, returning a list of axes
    and the keepdim flag. Handles both list-based dims (e.g., sum.dim_IntList) and
    single int dims (e.g., prod.dim_int).

    Args:
        args: The handler args list
        start_idx: Index where the dim argument starts (default 1, after self)

    Returns:
        Tuple of (axes, keepdim) where:
        - axes: List of dimension indices, or empty list for reduce-all
        - keepdim: Boolean keepdim flag (default False)
    """
    if len(args) > start_idx and isinstance(args[start_idx], (list, tuple)):
        dim = list(args[start_idx])
        keepdim = args[start_idx + 1] if len(args) > start_idx + 1 else False
    elif len(args) > start_idx and isinstance(args[start_idx], int):
        dim = [args[start_idx]]
        keepdim = args[start_idx + 1] if len(args) > start_idx + 1 else False
    else:
        dim = []
        keepdim = False

    return dim, keepdim


_UNARY_OPS: List[Tuple[Any, Any, str]] = [