Python · Elementwise and Reductions

Elementwise and Reductions

import numkong 8.0.0-rc.1 · b1f866ab

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.

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.

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.

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.

API

Functions

momentsCompute the sum and the sum of squares of all elements, returned as a pair.
__init__.pyi:1525 ↗
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]

Parameters

capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
minmaxFind the smallest and largest elements with their indices.
__init__.pyi:1535 ↗
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.

Parameters

capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
sumReturn the sum of all elements, or per-slice sums along one or more axes.
__init__.pyi:1548 ↗
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

Parameters

a_BufferTypeInput Tensor, NumPy array or any buffer-protocol object.
axisint | tuple[int, ...] | NoneAxis or axes to reduce along, or None to reduce every element to a scalar.
keepdimsboolKeep each reduced axis as a size-1 dimension.
outTensor | NonePreallocated output Tensor for the result.
dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneInput dtype to read a as, overriding the one it reports.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
normReturn the L2 norm of all elements, or per-slice norms along one or more axes.
__init__.pyi:1569 ↗
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

Parameters

a_BufferTypeInput Tensor, NumPy array or any buffer-protocol object.
axisint | tuple[int, ...] | NoneAxis or axes to reduce along, or None to reduce every element to a scalar.
keepdimsboolKeep each reduced axis as a size-1 dimension.
outTensor | NonePreallocated output Tensor for the result.
dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneInput dtype to read a as, overriding the one it reports.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
minReturn the minimum element, or per-slice minimums along one or more axes.
__init__.pyi:1590 ↗
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.

Parameters

a_BufferTypeInput Tensor, NumPy array or any buffer-protocol object.
axisint | tuple[int, ...] | NoneAxis or axes to reduce along, or None to reduce every element to a scalar.
keepdimsboolKeep each reduced axis as a size-1 dimension.
outTensor | NonePreallocated output Tensor for the result.
dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneInput dtype to read a as, overriding the one it reports.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
maxReturn the maximum element, or per-slice maximums along one or more axes.
__init__.pyi:1613 ↗
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.

Parameters

a_BufferTypeInput Tensor, NumPy array or any buffer-protocol object.
axisint | tuple[int, ...] | NoneAxis or axes to reduce along, or None to reduce every element to a scalar.
keepdimsboolKeep each reduced axis as a size-1 dimension.
outTensor | NonePreallocated output Tensor for the result.
dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneInput dtype to read a as, overriding the one it reports.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
argminReturn the index of the minimum element, or per-slice indices along one axis.
__init__.pyi:1636 ↗
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.

Parameters

a_BufferTypeInput Tensor, NumPy array or any buffer-protocol object.
axisint | NoneAxis to reduce along, or None to reduce every element to a scalar.
keepdimsboolKeep each reduced axis as a size-1 dimension.
outTensor | NonePreallocated output Tensor for the result.
dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneInput dtype to read a as, overriding the one it reports.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
argmaxReturn the index of the maximum element, or per-slice indices along one axis.
__init__.pyi:1659 ↗
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.

Parameters

a_BufferTypeInput Tensor, NumPy array or any buffer-protocol object.
axisint | NoneAxis to reduce along, or None to reduce every element to a scalar.
keepdimsboolKeep each reduced axis as a size-1 dimension.
outTensor | NonePreallocated output Tensor for the result.
dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneInput dtype to read a as, overriding the one it reports.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
fmaCompute alpha × a × b + beta × c element-wise over three tensors of one shape.
__init__.pyi:1686 ↗
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.

Parameters

a_BufferTypeFirst tensor, of any rank.
b_BufferTypeSecond tensor, of the shape of a.
c_BufferTypeThird tensor, of the shape of a.
dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneInput dtype to read every buffer as, overriding the one they report.
alphafloatScale of the product of a and b.
betafloatScale of c.
out_BufferType | NoneOutput tensor for the result.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
blendCompute alpha × a + beta × b element-wise over two tensors of one shape.
__init__.pyi:1713 ↗
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.

Parameters

a_BufferTypeFirst tensor, of any rank.
b_BufferTypeSecond tensor, of the shape of a.
dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneInput dtype to read both buffers as, overriding the one they report.
alphafloatScale of a.
betafloatScale of b.
out_BufferType | NoneOutput tensor for the result.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
sinCompute the sine of every element.
__init__.pyi:1742 ↗
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.

Parameters

a_BufferTypeInput tensor of any rank, angles in radians.
dtype_FloatTypeName | _MiniFloatType | NoneInput dtype to read a as, overriding the one it reports.
out_BufferType | NoneOutput tensor for the result.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
cosCompute the cosine of every element.
__init__.pyi:1761 ↗
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.

Parameters

a_BufferTypeInput tensor of any rank, angles in radians.
dtype_FloatTypeName | _MiniFloatType | NoneInput dtype to read a as, overriding the one it reports.
out_BufferType | NoneOutput tensor for the result.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
atanCompute the arctangent of every element, in radians.
__init__.pyi:1780 ↗
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.

Parameters

a_BufferTypeInput tensor of any rank.
dtype_FloatTypeName | _MiniFloatType | NoneInput dtype to read a as, overriding the one it reports.
out_BufferType | NoneOutput tensor for the result.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
scaleCompute alpha × a + beta element-wise over a tensor of any rank.
__init__.pyi:1803 ↗
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.

Parameters

a_BufferTypeInput tensor of any rank.
dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneInput dtype to read a as, overriding the one it reports.
alphafloatMultiplicative scale.
betafloatAdditive offset.
out_BufferType | NoneOutput tensor for the result.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
addAdd two tensors, or a tensor and a scalar, element-wise with NumPy broadcasting.
__init__.pyi:1826 ↗
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.

Parameters

a_BufferType | float | intFirst operand, a tensor or a scalar.
b_BufferType | float | intSecond operand, a tensor or a scalar.
out_BufferType | NoneOutput buffer for the result.
a_dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneOverride of the dtype of a.
b_dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneOverride of the dtype of b.
out_dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneOverride of the dtype of the result.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
multiplyMultiply two tensors, or a tensor and a scalar, element-wise with NumPy broadcasting.
__init__.pyi:1851 ↗
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.

Parameters

a_BufferType | float | intFirst operand, a tensor or a scalar.
b_BufferType | float | intSecond operand, a tensor or a scalar.
out_BufferType | NoneOutput buffer for the result.
a_dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneOverride of the dtype of a.
b_dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneOverride of the dtype of b.
out_dtype_FloatTypeName | Literal['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _MiniFloatType | NoneOverride of the dtype of the result.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
rmsnormGrouped RMSNorm, y = x * rsqrt(mean(x^2) + epsilon) * gamma, in place unless out is given.
__init__.pyi:1880 ↗
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.

Parameters

x_BufferTypeInput of dtype float32, bfloat16, or e4m3, with a contiguous last axis.
gamma_BufferType | NonePer-column float32 gain of length columns, defaulting to unit scale.
out_BufferType | NoneOutput buffer, same shape and dtype as x, may alias x.
groupsintIndependent sub-vectors per row.
epsilonfloatVariance epsilon. For e4m3, pass it in code units and fold the output scale into gamma.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
rmscastGrouped RMSNorm of dot products back into the dtype they multiplied, rounding like the cast.
__init__.pyi:1904 ↗
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.

Parameters

x_BufferTypeDot 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_BufferType | NonePer-column float32 gain of length columns, defaulting to unit scale.
dtypeLiteral['uint1', 'int4', 'uint4', 'int64', 'int32', 'int16', 'int8', 'uint64', 'uint32', 'uint16', 'uint8'] | _FloatTypeName | _MiniFloatTypeOutput dtype, which names the dtype x must have.
out_BufferType | NoneOutput buffer of the shape of x and of dtype.
groupsintIndependent sub-vectors per row.
epsilonfloatVariance epsilon, added to the mean square of the raw input.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.
swigluFused SwiGLU, y = silu(gate) * up, in place unless out is given.
__init__.pyi:1934 ↗
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.

Parameters

gate_BufferTypeGate input of dtype float32, bfloat16, or e4m3, with a contiguous last axis.
up_BufferType | NoneUp input, same shape and dtype as gate, defaulting to plain SiLU.
out_BufferType | NoneOutput buffer, same shape and dtype as gate, may alias gate.
gate_scalefloatTensor scale of gate.
output_scalefloatMultiplier on the result.
capabilitiesCapability | int | NoneCapabilities to run within: every CPU one by default, or for packed operands the mask that packed them.
streamint | NoneA GPU stream pointer, left None on the CPU.