Reduction Operators

本页暂无中文版。以下为英文原文。

Every op on this page is used the same way: construct it once, then call it. The constructor takes what the kernel is compiled with; the call takes the tensors. Both are documented under each op — __init__ and forward, where forward is what runs when you call op(...).

Sums, means and extrema

tileops.ops.reduction.reduce.SumFwdOp

Sum reduction along dim=-1.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Construct a reduce op.

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], tuple[int, ...], or None.

• keepdim (bool, default: False ) –

  Whether to retain reduced dims as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.reduce.MeanFwdOp

Mean reduction along dim=-1.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Construct a reduce op.

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], tuple[int, ...], or None.

• keepdim (bool, default: False ) –

  Whether to retain reduced dims as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.reduce.ProdFwdOp

Product reduction.

Unlike the other simple reduce ops, ProdFwdOp defaults to dim=-1 (manifest declares default: -1 for prod).

__init__

__init__(
    dim=-1,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Construct ProdFwdOp.

Parameters:

• dim (int, default: -1 ) –

  Reduction dimension (default -1).

• keepdim (bool, default: False ) –

  Whether to retain reduced dims as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.reduce.AmaxFwdOp

Amax (element-wise maximum) reduction along dim=-1.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Construct a reduce op.

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], tuple[int, ...], or None.

• keepdim (bool, default: False ) –

  Whether to retain reduced dims as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.reduce.AminFwdOp

Amin (element-wise minimum) reduction along dim=-1.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Construct a reduce op.

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], tuple[int, ...], or None.

• keepdim (bool, default: False ) –

  Whether to retain reduced dims as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

Variance and deviation

tileops.ops.reduction.reduce.VarFwdOp

Variance reduction with Bessel's correction.

__init__

__init__(
    dim=None,
    correction=1,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Construct a Welford-based reduce op.

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], tuple[int, ...], or None.

• correction (int, default: 1 ) –

  Bessel's correction (default 1).

• keepdim (bool, default: False ) –

  Whether to retain reduced dims as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], or tuple[int, ...] for multi-dim reduction.

• correction (int, default: 1 ) –

  Bessel's correction (default 1).

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.reduce.VarMeanFwdOp

Variance and mean reduction.

__init__

__init__(
    dim=None,
    correction=1,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Construct a Welford-based reduce op.

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], tuple[int, ...], or None.

• correction (int, default: 1 ) –

  Bessel's correction (default 1).

• keepdim (bool, default: False ) –

  Whether to retain reduced dims as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], or tuple[int, ...] for multi-dim reduction.

• correction (int, default: 1 ) –

  Bessel's correction (default 1).

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.reduce.StdFwdOp

Standard deviation reduction with Bessel's correction.

__init__

__init__(
    dim=None,
    correction=1,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Construct a Welford-based reduce op.

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], tuple[int, ...], or None.

• correction (int, default: 1 ) –

  Bessel's correction (default 1).

• keepdim (bool, default: False ) –

  Whether to retain reduced dims as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], or tuple[int, ...] for multi-dim reduction.

• correction (int, default: 1 ) –

  Bessel's correction (default 1).

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

Arg reductions

tileops.ops.reduction.argreduce.ArgmaxFwdOp

Argmax reduction along an arbitrary dim, returning int64 indices.

Construction: ArgmaxFwdOp(dim=None, keepdim=False).

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (Optional[int], default: None ) –

  Reduction dimension. None (the default) matches torch.argmax(x) semantics: the input is treated as a contiguous flattened 1D buffer and the returned index is into that flattened tensor.

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional custom kernel map.

• tune (bool, default: False ) –

  Whether to autotune the kernel.

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.argreduce.ArgminFwdOp

Argmin reduction along an arbitrary dim, returning int64 indices.

Construction: ArgminFwdOp(dim=None, keepdim=False).

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (Optional[int], default: None ) –

  Reduction dimension. None (the default) matches torch.argmin(x) semantics: the input is treated as a contiguous flattened 1D buffer and the returned index is into that flattened tensor.

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional custom kernel map.

• tune (bool, default: False ) –

  Whether to autotune the kernel.

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

Softmax

tileops.ops.reduction.softmax.SoftmaxFwdOp

Softmax operator: y = softmax(x, dim).

Output has the same shape and dtype as input. The reduction-dim extent N and dtype are inferred from x during forward().

__init__

__init__(
    dim=None,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (Optional[int], default: None ) –

  Reduction dimension (default None, matching PyTorch's torch.nn.functional.softmax). When None, the axis is resolved at forward time using PyTorch's implicit-axis rule (0 for ndim in {0, 1, 3} else 1) and the same deprecation UserWarning is emitted.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the softmax-family op.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.softmax.LogSoftmaxFwdOp

Log-softmax operator: y = log_softmax(x, dim).

Output has the same shape and dtype as input. The reduction-dim extent N and dtype are inferred from x during forward().

__init__

__init__(
    dim=None,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (Optional[int], default: None ) –

  Reduction dimension (default None, matching PyTorch's torch.nn.functional.log_softmax). When None, the axis is resolved at forward time using PyTorch's implicit-axis rule (0 for ndim in {0, 1, 3} else 1) and the same deprecation UserWarning is emitted.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the softmax-family op.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.softmax.LogSumExpFwdOp

LogSumExp operator: y = logsumexp(x, dim, keepdim).

Output shape is input shape without the reduction dimension (or with size-1 if keepdim=True).

__init__

__init__(
    dim=-1,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (Union[int, List[int]], default: -1 ) –

  Reduction dimension (default -1).

• keepdim (bool, default: False ) –

  Retain reduced dimension (default False).

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

forward

forward(
    x,
)

Run the softmax-family op.

One call to the operator this op registers: this is as far as dynamo traces.

Vector norms

tileops.ops.reduction.vector_norm.L1NormFwdOp

L1 norm reduction along a configurable dim.

Construction: L1NormFwdOp(dim=None, keepdim=False).

__init__

__init__(
    ord=1,
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (Union[int, List[int], None], default: None ) –

  Reduction dimension (default None -> full reduction, matching torch.linalg.vector_norm). Accepts int, list[int], or None.

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• ord (Union[int, float], default: 1 ) –

  Norm order. Must equal 1 for L1NormFwdOp (manifest fixes ord == 1); accepted as a kwarg to mirror torch.linalg.vector_norm.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional custom kernel map.

• tune (bool, default: False ) –

  Whether to autotune the kernel.

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.vector_norm.L2NormFwdOp

L2 norm reduction along a configurable dim.

Construction: L2NormFwdOp(dim=None, keepdim=False).

__init__

__init__(
    ord=2,
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (Union[int, List[int], None], default: None ) –

  Reduction dimension (default None -> full reduction, matching torch.linalg.vector_norm). Accepts int, list[int], or None.

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• ord (Union[int, float], default: 2 ) –

  Norm order. Must equal 2 for L2NormFwdOp (manifest fixes ord == 2); accepted as a kwarg to mirror torch.linalg.vector_norm.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional custom kernel map.

• tune (bool, default: False ) –

  Whether to autotune the kernel.

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.vector_norm.InfNormFwdOp

Infinity norm reduction along a configurable dim.

Construction: InfNormFwdOp(dim=None, keepdim=False).

NaN handling: rows containing any NaN produce NaN output, matching torch.linalg.vector_norm(ord=inf) semantics. The kernel drops NaN values and patches those rows itself, so the compensation stays with the implementation that needs it.

__init__

__init__(
    ord=inf,
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (Union[int, List[int], None], default: None ) –

  Reduction dimension (default None -> full reduction, matching torch.linalg.vector_norm). Accepts int, list[int], or None.

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• ord (Union[int, float], default: inf ) –

  Norm order. Must equal float('inf') for InfNormFwdOp (manifest fixes ord == float('inf')); accepted as a kwarg to mirror torch.linalg.vector_norm.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional custom kernel map.

• tune (bool, default: False ) –

  Whether to autotune the kernel.

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

Cumulative

tileops.ops.reduction.cumulative.CumulativeOp

Abstract base for cumulative scan operators with a user-selectable axis.

Subclasses must override _op_kind (class attribute) — the kernel's op-kind dispatch string ("sum" or "prod").

__init__

__init__(
    dim=-1,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (int, default: -1 ) –

  Reduction axis (default -1). Negative values are normalized at forward time (dim % x.ndim).

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  If True, autotune tile configs.

forward

forward(
    x,
)

Run the scan.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.cumulative.CumsumFwdOp

Cumulative sum operator: y = cumsum(x, dim).

Output has the same shape and dtype as x. Alignment padding is handled inside the kernel via masked loads.

Shapes with M < 128 and N > 8192 take a three-pass parallel scan for SM utilization; every other shape takes the sequential scan.

Parameters:

• dim (int, default: -1 ) –

  Reduction axis (default -1). Negative values are normalized at forward time.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

Example

op = CumsumFwdOp()
x = torch.randn(1024, 4096, dtype=torch.float16, device="cuda")
y = op(x)  # shape: (1024, 4096)

__init__

__init__(
    dim=-1,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (int, default: -1 ) –

  Reduction axis (default -1). Negative values are normalized at forward time (dim % x.ndim).

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  If True, autotune tile configs.

forward

forward(
    x,
)

Run the scan.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.cumulative.CumprodFwdOp

Cumulative product operator: y = cumprod(x, dim).

Output has the same shape and dtype as x. Alignment padding is handled inside the kernel via masked loads.

Parameters:

• dim (int, default: -1 ) –

  Reduction axis (default -1). Negative values are normalized at forward time.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

Example

op = CumprodFwdOp()
x = torch.randn(1024, 4096, dtype=torch.float16, device="cuda") * 0.01 + 0.99
y = op(x)  # shape: (1024, 4096)

__init__

__init__(
    dim=-1,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (int, default: -1 ) –

  Reduction axis (default -1). Negative values are normalized at forward time (dim % x.ndim).

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional kernel override dict.

• tune (bool, default: False ) –

  If True, autotune tile configs.

forward

forward(
    x,
)

Run the scan.

One call to the operator this op registers: this is as far as dynamo traces.

Logical reductions

tileops.ops.reduction.logical_reduce.AllFwdOp

All reduction along dim, returning bool.

Construction: AllFwdOp(dim=None, keepdim=False). are derived from the input tensor at forward time, and kernels are cached by (M, N) to avoid rebuilds.

Supports any numeric dtype including torch.bool, int32, int64, and complex types. A dtype TileLang cannot store as shared memory is converted inside the kernel, so this op hands over the tensor its manifest declares.

Empty-dim contract: dim=[] / dim=() is a no-op -- forward returns x.bool() with the input shape, matching torch.all semantics.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Construct AllFwdOp.

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], tuple[int, ...], or None.

• keepdim (bool, default: False ) –

  Whether to retain reduced dims as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], or tuple[int, ...] for multi-dim reduction.

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional custom kernel map.

• tune (bool, default: False ) –

  Whether to autotune the kernel.

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.logical_reduce.AnyFwdOp

Any reduction along dim, returning bool.

Construction: AnyFwdOp(dim=None, keepdim=False). are derived from the input tensor at forward time, and kernels are cached by (M, N) to avoid rebuilds.

Supports any numeric dtype including torch.bool, int32, int64, and complex types. A dtype TileLang cannot store as shared memory is converted inside the kernel, so this op hands over the tensor its manifest declares.

Empty-dim contract: dim=[] / dim=() is a no-op -- forward returns x.bool() with the input shape, matching torch.any semantics.

__init__

__init__(
    dim=None,
    keepdim=False,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Construct AnyFwdOp.

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], tuple[int, ...], or None.

• keepdim (bool, default: False ) –

  Whether to retain reduced dims as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional override for kernel dispatch.

• tune (bool, default: False ) –

  Whether to autotune (default False).

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], or tuple[int, ...] for multi-dim reduction.

• keepdim (bool, default: False ) –

  Whether to retain the reduced dimension as size 1.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional custom kernel map.

• tune (bool, default: False ) –

  Whether to autotune the kernel.

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.

tileops.ops.reduction.logical_reduce.CountNonzeroFwdOp

Count nonzero reduction along dim, returning int64.

Construction: CountNonzeroFwdOp(dim=None).

Note: No keepdim parameter -- the reduction dimension is always removed, matching torch.count_nonzero semantics.

Supports any numeric dtype including torch.bool, int32, int64, and complex types. A dtype TileLang cannot store as shared memory is converted inside the kernel, so this op hands over the tensor its manifest declares.

__init__

__init__(
    dim=None,
    *,
    target=None,
    kernel_map=None,
    tune=False
)

Build the op. Shapes and dtype are taken from the first call.

Parameters:

• dim (Union[int, List[int], Tuple[int, ...], None], default: None ) –

  Reduction dimension (default None, i.e. full reduction). Accepts int, list[int], or tuple[int, ...] for multi-dim reduction.

• target (Target, default: None ) –

  Which set of kernels serves this op — a target name, BUILTIN for the in-tree kernels, or None to decide from the input device.

• kernel_map (Optional[Dict[str, Kernel]], default: None ) –

  Optional custom kernel map.

• tune (bool, default: False ) –

  Whether to autotune the kernel.

forward

forward(
    x,
)

Run the reduce op on x along the configured dim.

One call to the operator this op registers: this is as far as dynamo traces.