# Elementwise and Reductions

> NumKong for Python

Elementwise arithmetic and fused operations are their own family.
They share the tensor infrastructure but should not be collapsed into the reduction or matrix sections.

```python
import numpy as np
import numkong as nk

a = np.arange(8, dtype=np.float32)
b = np.arange(8, dtype=np.float32)[::-1].copy()

scaled = nk.scale(a, alpha=2.0, beta=1.0)     # 2 * a + 1
blended = nk.blend(a, b, alpha=0.25, beta=0.75)
fused = nk.fma(a, b, a, alpha=1.0, beta=1.0)  # a * b + a

assert np.asarray(scaled).shape == (8,)
assert np.asarray(fused).shape == (8,)
```

### Moments Reductions

Moments reductions return `(sum, sum_of_squares)`.
The key property is that NumKong does not force you into same-storage accumulation.

```python
import numpy as np
import numkong as nk

x = np.full(4096, 255, dtype=np.uint8)

nk_sum, nk_sumsq = nk.moments(nk.Tensor(x))
naive_sum = np.sum(x, dtype=np.uint8)      # overflows immediately
naive_sumsq = np.sum(x * x, dtype=np.uint8) # also overflows

print(nk_sum, nk_sumsq, naive_sum, naive_sumsq)
assert nk_sum > int(naive_sum)
assert nk_sumsq > int(naive_sumsq)
```

Same-width accumulation is a bad default for low-precision storage.

### Min/Max Reductions

Min/max reductions are in a separate section because they cover strided reduction cases.
NumKong provides SIMD-accelerated strided reductions that are not common in other libraries.

```python
import numpy as np
import numkong as nk

matrix = nk.Tensor(np.array([
    [ 3.0,  0.0, 7.0],
    [ 1.0,  2.0, 5.0],
    [ 4.0, -1.0, 6.0],
], dtype=np.float32))

second_column = matrix[:, 1]  # strided view into a row-major Nx3 tensor

idx = second_column.argmin()
mn, i0, mx, i1 = second_column.minmax()

assert idx == 2
assert int(i0) == 2
assert float(np.asarray(mn)) == -1.0
```

Fresh measurement for the rewritten docs:
on an Apple M2 Pro, `np.argmin(matrix[:, 1])` on a row-major `2,000,000 x 3` `float32` array took about `1.63 ms` median.
The equivalent NumKong `Tensor(... )[:, 1].argmin()` took about `0.67 ms` median.
That is about `2.45x` faster on this strided reduction case.

`import numkong`

## moments

Compute the sum and the sum of squares of all elements, returned as a pair.

```python
def moments(
    a: _BufferType,
    /,
    *,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> tuple[float, float]
```

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## minmax

Find the smallest and largest elements with their indices.

```python
def minmax(
    a: _BufferType,
    /,
    *,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> tuple[float, int, float, int] | None
```

Returns `(min_value, min_index, max_value, max_index)`, or None when every element is NaN.

- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## sum

Return the sum of all elements, or per-slice sums along one or more axes.

```python
def sum(
    a: _BufferType,
    /,
    axis: int | tuple[int, ...] | None = None,
    *,
    keepdims: bool = False,
    out: Tensor | None = None,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> float | int | Tensor
```

- `a`: Input Tensor, NumPy array or any buffer-protocol object.
- `axis`: Axis or axes to reduce along, or None to reduce every element to a scalar.
- `keepdims`: Keep each reduced axis as a size-1 dimension.
- `out`: Preallocated output Tensor for the result.
- `dtype`: Input dtype to read `a` as, overriding the one it reports.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## norm

Return the L2 norm of all elements, or per-slice norms along one or more axes.

```python
def norm(
    a: _BufferType,
    /,
    axis: int | tuple[int, ...] | None = None,
    *,
    keepdims: bool = False,
    out: Tensor | None = None,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> float | Tensor
```

- `a`: Input Tensor, NumPy array or any buffer-protocol object.
- `axis`: Axis or axes to reduce along, or None to reduce every element to a scalar.
- `keepdims`: Keep each reduced axis as a size-1 dimension.
- `out`: Preallocated output Tensor for the result.
- `dtype`: Input dtype to read `a` as, overriding the one it reports.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## min

Return the minimum element, or per-slice minimums along one or more axes.

```python
def min(
    a: _BufferType,
    /,
    axis: int | tuple[int, ...] | None = None,
    *,
    keepdims: bool = False,
    out: Tensor | None = None,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> float | int | Tensor | None
```

The scalar result is None when every element is NaN.

- `a`: Input Tensor, NumPy array or any buffer-protocol object.
- `axis`: Axis or axes to reduce along, or None to reduce every element to a scalar.
- `keepdims`: Keep each reduced axis as a size-1 dimension.
- `out`: Preallocated output Tensor for the result.
- `dtype`: Input dtype to read `a` as, overriding the one it reports.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## max

Return the maximum element, or per-slice maximums along one or more axes.

```python
def max(
    a: _BufferType,
    /,
    axis: int | tuple[int, ...] | None = None,
    *,
    keepdims: bool = False,
    out: Tensor | None = None,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> float | int | Tensor | None
```

The scalar result is None when every element is NaN.

- `a`: Input Tensor, NumPy array or any buffer-protocol object.
- `axis`: Axis or axes to reduce along, or None to reduce every element to a scalar.
- `keepdims`: Keep each reduced axis as a size-1 dimension.
- `out`: Preallocated output Tensor for the result.
- `dtype`: Input dtype to read `a` as, overriding the one it reports.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## argmin

Return the index of the minimum element, or per-slice indices along one axis.

```python
def argmin(
    a: _BufferType,
    /,
    axis: int | None = None,
    *,
    keepdims: bool = False,
    out: Tensor | None = None,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> int | Tensor | None
```

The scalar result is None when every element is NaN.

- `a`: Input Tensor, NumPy array or any buffer-protocol object.
- `axis`: Axis to reduce along, or None to reduce every element to a scalar.
- `keepdims`: Keep each reduced axis as a size-1 dimension.
- `out`: Preallocated output Tensor for the result.
- `dtype`: Input dtype to read `a` as, overriding the one it reports.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## argmax

Return the index of the maximum element, or per-slice indices along one axis.

```python
def argmax(
    a: _BufferType,
    /,
    axis: int | None = None,
    *,
    keepdims: bool = False,
    out: Tensor | None = None,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> int | Tensor | None
```

The scalar result is None when every element is NaN.

- `a`: Input Tensor, NumPy array or any buffer-protocol object.
- `axis`: Axis to reduce along, or None to reduce every element to a scalar.
- `keepdims`: Keep each reduced axis as a size-1 dimension.
- `out`: Preallocated output Tensor for the result.
- `dtype`: Input dtype to read `a` as, overriding the one it reports.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## fma

Compute alpha × a × b + beta × c element-wise over three tensors of one shape.

```python
def fma(
    a: _BufferType,
    b: _BufferType,
    c: _BufferType,
    /,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    *,
    alpha: float = 1,
    beta: float = 1,
    out: _BufferType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

Returns a new Tensor, or None after writing into `out`.

- `a`: First tensor, of any rank.
- `b`: Second tensor, of the shape of `a`.
- `c`: Third tensor, of the shape of `a`.
- `dtype`: Input dtype to read every buffer as, overriding the one they report.
- `alpha`: Scale of the product of `a` and `b`.
- `beta`: Scale of `c`.
- `out`: Output tensor for the result.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## blend

Compute alpha × a + beta × b element-wise over two tensors of one shape.

```python
def blend(
    a: _BufferType,
    b: _BufferType,
    /,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    *,
    alpha: float = 1,
    beta: float = 1,
    out: _BufferType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

Returns a new Tensor, or None after writing into `out`.

- `a`: First tensor, of any rank.
- `b`: Second tensor, of the shape of `a`.
- `dtype`: Input dtype to read both buffers as, overriding the one they report.
- `alpha`: Scale of `a`.
- `beta`: Scale of `b`.
- `out`: Output tensor for the result.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## sin

Compute the sine of every element.

```python
def sin(
    a: _BufferType,
    /,
    dtype: _FloatTypeName | _MiniFloatType | None = None,
    *,
    out: _BufferType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

Returns a new Tensor, or None after writing into `out`.

- `a`: Input tensor of any rank, angles in radians.
- `dtype`: Input dtype to read `a` as, overriding the one it reports.
- `out`: Output tensor for the result.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## cos

Compute the cosine of every element.

```python
def cos(
    a: _BufferType,
    /,
    dtype: _FloatTypeName | _MiniFloatType | None = None,
    *,
    out: _BufferType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

Returns a new Tensor, or None after writing into `out`.

- `a`: Input tensor of any rank, angles in radians.
- `dtype`: Input dtype to read `a` as, overriding the one it reports.
- `out`: Output tensor for the result.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## atan

Compute the arctangent of every element, in radians.

```python
def atan(
    a: _BufferType,
    /,
    dtype: _FloatTypeName | _MiniFloatType | None = None,
    *,
    out: _BufferType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

Returns a new Tensor, or None after writing into `out`.

- `a`: Input tensor of any rank.
- `dtype`: Input dtype to read `a` as, overriding the one it reports.
- `out`: Output tensor for the result.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## scale

Compute alpha × a + beta element-wise over a tensor of any rank.

```python
def scale(
    a: _BufferType,
    /,
    dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    *,
    alpha: float = 1,
    beta: float = 0,
    out: _BufferType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

Returns a new Tensor, or None after writing into `out`.

- `a`: Input tensor of any rank.
- `dtype`: Input dtype to read `a` as, overriding the one it reports.
- `alpha`: Multiplicative scale.
- `beta`: Additive offset.
- `out`: Output tensor for the result.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## add

Add two tensors, or a tensor and a scalar, element-wise with NumPy broadcasting.

```python
def add(
    a: _BufferType | float | int,
    b: _BufferType | float | int,
    /,
    *,
    out: _BufferType | None = None,
    a_dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    b_dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    out_dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

Equivalent to `a + b`. Returns a new Tensor, or None after writing into `out`.

- `a`: First operand, a tensor or a scalar.
- `b`: Second operand, a tensor or a scalar.
- `out`: Output buffer for the result.
- `a_dtype`: Override of the dtype of `a`.
- `b_dtype`: Override of the dtype of `b`.
- `out_dtype`: Override of the dtype of the result.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## multiply

Multiply two tensors, or a tensor and a scalar, element-wise with NumPy broadcasting.

```python
def multiply(
    a: _BufferType | float | int,
    b: _BufferType | float | int,
    /,
    *,
    out: _BufferType | None = None,
    a_dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    b_dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    out_dtype: _FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | None = None,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

Equivalent to `a * b`. Returns a new Tensor, or None after writing into `out`.

- `a`: First operand, a tensor or a scalar.
- `b`: Second operand, a tensor or a scalar.
- `out`: Output buffer for the result.
- `a_dtype`: Override of the dtype of `a`.
- `b_dtype`: Override of the dtype of `b`.
- `out_dtype`: Override of the dtype of the result.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## rmsnorm

Grouped RMSNorm, `y = x * rsqrt(mean(x^2) + epsilon) * gamma`, in place unless `out` is given.

```python
def rmsnorm(
    x: _BufferType,
    /,
    gamma: _BufferType | None = None,
    *,
    out: _BufferType | None = None,
    groups: int = 1,
    epsilon: float = 1e-06,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

Each row, spanning all axes but the last, holds `groups` independent sub-vectors of length
`x.shape[-1] // groups`, normalized separately.

- `x`: Input of dtype float32, bfloat16, or e4m3, with a contiguous last axis.
- `gamma`: Per-column float32 gain of length [`columns`](https://www.unum.cloud/docs/numkong/python/matrix-kernels.md#columns), defaulting to unit scale.
- `out`: Output buffer, same shape and dtype as `x`, may alias `x`.
- `groups`: Independent sub-vectors per row.
- `epsilon`: Variance epsilon. For e4m3, pass it in code units and fold the output scale into `gamma`.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## rmscast

Grouped RMSNorm of dot products back into the [`dtype`](https://www.unum.cloud/docs/numkong/python/mesh-alignment.md#dtype) they multiplied, rounding like the cast.

```python
def rmscast(
    x: _BufferType,
    /,
    gamma: _BufferType | None = None,
    *,
    dtype: Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _MiniFloatType,
    out: _BufferType | None = None,
    groups: int = 1,
    epsilon: float = 1e-06,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

Each row, spanning all axes but the last, holds `groups` independent sub-vectors of length
`x.shape[-1] // groups`, normalized separately.

Integer outputs round to nearest even and saturate, with `gamma` carrying the quantization
scale; uint8 suits only non-negative activations.

- `x`: Dot products with a contiguous last axis: float32 for bfloat16, float16, e4m3, e5m2, e2m3
and e3m2 outputs, float64 for float32, int32 for int8 and uint32 for uint8.
- `gamma`: Per-column float32 gain of length [`columns`](https://www.unum.cloud/docs/numkong/python/matrix-kernels.md#columns), defaulting to unit scale.
- `dtype`: Output dtype, which names the dtype `x` must have.
- `out`: Output buffer of the shape of `x` and of [`dtype`](https://www.unum.cloud/docs/numkong/python/mesh-alignment.md#dtype).
- `groups`: Independent sub-vectors per row.
- `epsilon`: Variance epsilon, added to the mean square of the raw input.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.

## swiglu

Fused SwiGLU, `y = silu(gate) * up`, in place unless `out` is given.

```python
def swiglu(
    gate: _BufferType,
    /,
    up: _BufferType | None = None,
    *,
    out: _BufferType | None = None,
    gate_scale: float = 1.0,
    output_scale: float = 1.0,
    capabilities: Capability | int | None = ...,
    stream: int | None = ...,
) -> Tensor | None
```

The gate is read at `gate_scale * gate` and the result is multiplied by `output_scale`.
With `up` omitted this reduces to plain SiLU.

- `gate`: Gate input of dtype float32, bfloat16, or e4m3, with a contiguous last axis.
- `up`: Up input, same shape and dtype as `gate`, defaulting to plain SiLU.
- `out`: Output buffer, same shape and dtype as `gate`, may alias `gate`.
- `gate_scale`: Tensor scale of `gate`.
- `output_scale`: Multiplier on the result.
- `capabilities`: Capabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
- `stream`: A GPU stream pointer, left None on the CPU.