# Source code for tensorplay.nn.functional

Source: https://www.tensorplay.cn/docs/_modules/tensorplay/nn/functional.html

```
"""Functional interface."""

import enum
import math
import os
import warnings
from typing import Any, Optional

import tensorplay
import tensorplay._C as _C
from tensorplay._C import _add_docstr, DType
from tensorplay import Tensor
from tensorplay.graph import capture_call as _capture_call

[docs]
def threshold(
    input: Tensor,
    threshold: float,
    value: float,
    inplace: bool = False,
) -> Tensor:
    r"""Apply a threshold to each element of the input Tensor.

    See :class:`~tensorplay.nn.Threshold` for more details.
    """
    captured = _capture_call(
        globals()["threshold"], (input, threshold, value), {"inplace": inplace}
    )
    if captured is not None:
        return captured
    if inplace:
        result = _C.threshold_(input, threshold, value)
    else:
        result = _C.threshold(input, threshold, value)
    return result

[docs]
def silu(input: Tensor, inplace: bool = False) -> Tensor:
    r"""Apply the Sigmoid Linear Unit (SiLU) function, element-wise.

    The SiLU function is also known as the swish function.

    .. math::
        \text{silu}(x) = x * \sigma(x), \text{where } \sigma(x) \text{ is the logistic sigmoid.}

    .. note::
        See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_
        where the SiLU (Sigmoid Linear Unit) was originally coined, and see
        `Sigmoid-Weighted Linear Units for Neural Network Function Approximation
        in Reinforcement Learning <https://arxiv.org/abs/1702.03118>`_ and `Swish:
        a Self-Gated Activation Function <https://arxiv.org/abs/1710.05941v1>`_
        where the SiLU was experimented with later.

    See :class:`~tensorplay.nn.SiLU` for more details.
    """
    captured = _capture_call(silu, (input,), {"inplace": inplace})
    if captured is not None:
        return captured
    if inplace:
        return tensorplay._C.silu_(input)
    return tensorplay._C.silu(input)

gelu = _add_docstr(
    tensorplay._C.gelu,
    r"""
gelu(input, approximate = 'none') -> Tensor

When the approximate argument is 'none', it applies element-wise the function
:math:`\text{GELU}(x) = x * \Phi(x)`

where :math:`\Phi(x)` is the Cumulative Distribution Function for Gaussian Distribution.

When the approximate argument is 'tanh', Gelu is estimated with

.. math::
    \text{GELU}(x) = 0.5 * x * (1 + \text{Tanh}(\sqrt{2 / \pi} * (x + 0.044715 * x^3)))

See `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_.
""",
)

[docs]
def linear(input: Tensor, weight: Tensor, bias: Optional[Tensor] = None) -> Tensor:
    r"""Applies a linear transformation to the incoming data: :math:`y = xA^T + b`.

    Shape:
        - Input: :math:`(*, H_\text{in})` where :math:`*` means any number of
          dimensions including none and :math:`H_\text{in} = \text{in\_features}`.
        - Weight: :math:`(H_\text{out}, H_\text{in})` where
          :math:`H_\text{out} = \text{out\_features}`.
        - Bias: :math:`(H_\text{out})`
        - Output: :math:`(*, H_\text{out})`

    See :class:`~tensorplay.nn.Linear` for more details.
    """
    captured = _capture_call(linear, (input, weight, bias), {})
    if captured is not None:
        return captured

    distributed = _distributed_tensor_types()
    if distributed is not None and any(
        isinstance(value, distributed[0]) for value in (input, weight, bias)
    ):
        return _distributed_linear(input, weight, bias, distributed[0])

    # _matmul_impl checks this again later, but the native flatten path does
    # not work on scalar inputs, so try to catch this here already
    input_dim = input.dim()
    weight_dim = weight.dim()
    if input_dim == 0 or weight_dim == 0:
        raise RuntimeError(
            "both arguments to linear need to be at least 1D, but they are "
            f"{input_dim}D and {weight_dim}D"
        )

    # Native dispatch: CPU runs linear_kernel (single seeded-GEMM addmm with
    # raw as_strided weight.t(), bias folded into the epilogue / seed);
    # other backends fall through to the recordable matmul/add composite.
    return tensorplay.linear(input, weight, bias)

def _distributed_tensor_types() -> tuple[type, type] | None:
    try:
        from tensorplay.distributed.tensor import DTensor
    except (ImportError, AttributeError):
        return None
    return DTensor, Tensor

def _distributed_linear(input: Any, weight: Any, bias: Any, dtensor_type: type) -> Any:
    from tensorplay.distributed.tensor import DTensor, Partial, Shard

    input_dtensor = input if isinstance(input, dtensor_type) else None
    weight_dtensor = weight if isinstance(weight, dtensor_type) else None
    bias_value = bias.to_local() if isinstance(bias, dtensor_type) else bias
    input_value = input.to_local() if input_dtensor is not None else input
    weight_value = weight.to_local() if weight_dtensor is not None else weight
    result = tensorplay.linear(input_value, weight_value, bias_value)

    template = input_dtensor or weight_dtensor
    if template is None:
        return result
    placements = list(template.placements)
    if weight_dtensor is not None:
        for index, placement in enumerate(weight_dtensor.placements):
            if isinstance(placement, Shard) and placement.dim == 0:
                placements[index] = Shard(-1)
            elif isinstance(placement, Shard) and placement.dim == 1:
                placements[index] = Partial("sum")
    elif input_dtensor is not None:
        for index, placement in enumerate(input_dtensor.placements):
            if isinstance(placement, Shard) and placement.dim == input_dtensor.ndim - 1:
                placements[index] = Partial("sum")
    shape = tuple(input_value.shape[:-1]) + (int(weight.shape[0]),)
    stride = [1] * len(shape)
    for index in range(len(shape) - 2, -1, -1):
        stride[index] = stride[index + 1] * int(shape[index + 1])
    return DTensor.from_local(
        result,
        template.device_mesh,
        placements,
        shape=shape,
        stride=tuple(stride),
        run_check=False,
    )

[docs]
def bilinear(input1, input2, weight, bias=None):
    if list(input1.shape)[:-1] != list(input2.shape)[:-1]:
        raise ValueError("input1 and input2 must have the same batch dimensions")
    
    out_features, in1_features, in2_features = weight.shape
    
    # w: (Out, H1, H2) -> (Out, H2, H1)
    # TensorPlay permute expects a sequence
    w = weight.permute([0, 2, 1])
    
    # w: (Out * H2, H1)
    w = w.reshape(-1, in1_features)
    
    # input1: (*, H1)
    # input1 @ w.T: (*, H1) @ (H1, Out * H2) -> (*, Out * H2)
    temp = input1.matmul(w.t())
    
    # temp: (*, Out * H2)
    # Reshape to (*, Out, H2)
    new_shape = list(input1.shape)[:-1] + [out_features, in2_features]
    temp = temp.view(new_shape)
    
    # input2: (*, H2)
    # unsqueeze to (*, H2, 1)
    input2_expanded = input2.unsqueeze(-1)
    
    # temp: (*, Out, H2)
    # result: (*, Out, 1)
    output = temp.matmul(input2_expanded)
    
    # squeeze
    output = output.squeeze(-1)
    
    if bias is not None:
        output = output + bias
        
    return output

[docs]
def relu(input, inplace=False):
    captured = _capture_call(
        relu,
        (input,),
        {"inplace": True} if inplace else {},
    )
    if captured is not None:
        return captured
    if inplace:
        return _C.relu_(input)
    return _C.relu(input)

[docs]
def softmax(input, dim=None, dtype=None):
    if dim is None:
        dim = -1
    if dtype is None:
        dtype = tensorplay.undefined
    return input.softmax(dim, dtype)

[docs]
def log_softmax(input, dim=None, dtype=None):
    captured = _capture_call(log_softmax, (input, dim, dtype), {})
    if captured is not None:
        return captured
    if dim is None:
        dim = -1
    if dtype is None:
        dtype = tensorplay.undefined
    return _C.log_softmax(input, dim, dtype)

[docs]
def prelu(input, weight):
    captured = _capture_call(prelu, (input, weight), {})
    if captured is not None:
        return captured
    # PReLU(x) = max(0, x) + weight * min(0, x), one weight per channel (the
    # input's second axis).  The operator takes the weights as a flat list and
    # lays them along that axis itself.
    if weight.numel() != 1:
        if input.dim() < 2:
             raise ValueError("Input must have at least 2 dimensions when num_parameters > 1")

        # Check if num_parameters matches channel dim (dim 1)
        if input.size(1) != weight.numel():
            raise ValueError(f"num_parameters {weight.numel()} does not match input channel size {input.size(1)}")

    return _C.prelu(input, weight)

[docs]
def flatten(input, start_dim=0, end_dim=-1):
    return input.flatten(start_dim, end_dim)

[docs]
def embedding(input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False):
    captured = _capture_call(embedding, (input, weight, padding_idx, max_norm, norm_type, scale_grad_by_freq, sparse), {})
    if captured is not None:
        return captured
    try:
        from tensorplay.distributed.tensor import DTensor, Shard
    except (ImportError, AttributeError):
        DTensor = ()
        Shard = ()
    if DTensor and (isinstance(input, DTensor) or isinstance(weight, DTensor)):
        input_dtensor = input if isinstance(input, DTensor) else None
        weight_dtensor = weight if isinstance(weight, DTensor) else None
        input_value = input.to_local() if input_dtensor is not None else input
        weight_value = weight.to_local() if weight_dtensor is not None else weight
        if max_norm is not None:
            raise NotImplementedError("embedding: max_norm is not supported")
        if padding_idx is None:
            normalized_padding_idx = -1
        else:
            if padding_idx < -weight.size(0) or padding_idx >= weight.size(0):
                raise AssertionError("Padding_idx must be within num_embeddings")
            normalized_padding_idx = padding_idx + weight.size(0) if padding_idx < 0 else padding_idx
        result = _C.embedding(
            weight_value,
            input_value,
            normalized_padding_idx,
            scale_grad_by_freq,
            sparse,
        )
        template = weight_dtensor or input_dtensor
        placements = list(template.placements)
        if input_dtensor is not None and weight_dtensor is None:
            placements = list(input_dtensor.placements)
        shape = tuple(input.shape) + (int(weight.shape[1]),)
        stride = [1] * len(shape)
        for index in range(len(shape) - 2, -1, -1):
            stride[index] = stride[index + 1] * int(shape[index + 1])
        return DTensor.from_local(
            result,
            template.device_mesh,
            placements,
            shape=shape,
            stride=tuple(stride),
            run_check=False,
        )
    if max_norm is not None:
        _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)
    if padding_idx is None:
        padding_idx = -1
    else:
        if padding_idx < -weight.size(0) or padding_idx >= weight.size(0):
            raise AssertionError('Padding_idx must be within num_embeddings')
        if padding_idx < 0:
            padding_idx += weight.size(0)
    return _C.embedding(weight, input, padding_idx, scale_grad_by_freq, sparse)

# Add more functionals as needed

[docs]
def dropout(input, p=0.5, training=True, inplace=False):
    captured = _capture_call(dropout, (input, p, training, inplace), {})
    if captured is not None:
        return captured
    if p < 0 or p > 1:
        raise ValueError("dropout probability has to be between 0 and 1, but got {}".format(p))
    if not training or p == 0:
        return input

    if p == 1:
        result = _C.zeros_like(input)
        return input.copy_(result) if inplace else result

    if inplace:
        # dropout_ mutates self and records the mask for backward).
        mask = (_C.rand(input.shape, device=input.device) > p).to(input.dtype)
        return input.mul_(mask).mul_(1.0 / (1.0 - p))

    return _C.dropout(input, p, training)

[docs]
def dropout2d(input, p=0.5, training=True, inplace=False):
    if p < 0 or p > 1:
        raise ValueError("dropout probability has to be between 0 and 1, but got {}".format(p))
    if not training or p == 0:
        return input
        
    # Input must be at least 2D (N, C, ...)
    if input.dim() < 2:
         raise ValueError("Feature dropout requires at least 2 dimensions")

    if p == 1:
        result = _C.zeros_like(input)
        return input.copy_(result) if inplace else result
         
    shape = list(input.shape)
    shape[2:] = [1] * (input.dim() - 2)
    
    mask = (_C.rand(shape, device=input.device) > p).to(input.dtype)
    scale = 1.0 / (1.0 - p)
    
    if inplace:
        return input.mul_(mask).mul_(scale)
    else:
        return input * mask * scale

[docs]
def dropout3d(input, p=0.5, training=True, inplace=False):
    return dropout2d(input, p, training, inplace)

[docs]
def alpha_dropout(input, p=0.5, training=True, inplace=False):
    if p < 0 or p > 1:
        raise ValueError("dropout probability has to be between 0 and 1, but got {}".format(p))
    if not training or p == 0:
        return input

    # Native fused forward plus generated backward through the saved mask.
    if p == 1:
        result = _C.zeros_like(input)
    else:
        result, _mask = _C.native_alpha_dropout(input, p)
    if inplace:
        return input.copy_(result)
    return result

[docs]
def feature_dropout(input, p=0.5, training=False, inplace=False):
    r"""Randomly zeroes entire channels (dim 1)."""
    if p < 0 or p > 1:
        raise ValueError("dropout probability has to be between 0 and 1, but got {}".format(p))
    if input.dim() < 2:
        raise RuntimeError(
            f"Feature dropout requires at least 2 dimensions in the input, "
            f"but got {input.dim()}")
    if not training or p == 0 or input.numel() == 0:
        return input

    if p == 1:
        result = _C.zeros_like(input)
        return input.copy_(result) if inplace else result

    result, _mask = _C.native_feature_dropout(input, p)
    if inplace:
        return input.copy_(result)
    return result

[docs]
def dropout_(input, p=0.5, training=True):
    r"""In-place version of :func:`dropout`."""
    return dropout(input, p=p, training=training, inplace=True)

[docs]
def rrelu_(input, lower=1.0 / 8, upper=1.0 / 3, training=False):
    r"""In-place version of :func:`rrelu`."""
    return rrelu(input, lower=lower, upper=upper, training=training,
                 inplace=True)

[docs]
def feature_dropout_(input, p=0.5, training=True):
    r"""In-place version of :func:`feature_dropout`."""
    return feature_dropout(input, p=p, training=training, inplace=True)

# Pooling helpers
def _pair(x):
    if isinstance(x, (int, float)):
        return (x, x)
    return tuple(x)

def _single(x):
    if isinstance(x, (int, float)):
        return (x,)
    return tuple(x)

def _triple(x):
    if isinstance(x, (int, float)):
        return (x, x, x)
    return tuple(x)

def _conv_padding_mode(input, padding, stride, dilation, weight, ndim):
    """A padding spelled as a word instead of a count of cells.

    ``"valid"`` pads nothing.  ``"same"`` keeps every output position, which
    only holds at unit stride, and asks the padding to cover the kernel's
    reach: that reach split evenly, with an odd cell going to the trailing
    side.  A per-axis count pads both sides alike, so the trailing excess is
    put on the input up front -- nothing on the leading side, the difference
    on the trailing side -- and the convolution keeps the leading count.
    """
    if padding == "valid":
        return input, [0] * ndim
    if padding != "same":
        raise ValueError(f"invalid padding string: {padding!r}")
    if any(s != 1 for s in stride):
        raise ValueError(
            "padding='same' is not supported for strided convolutions"
        )
    pairs = []
    for axis in range(ndim):
        reach = dilation[axis] * (weight.size(2 + axis) - 1)
        lead = reach // 2
        pairs.append((lead, reach - lead))
    if any(lead != trail for lead, trail in pairs):
        excess = []
        for lead, trail in reversed(pairs):
            excess.extend((0, trail - lead))
        input = pad(input, excess)
    return input, [lead for lead, _ in pairs]

[docs]
def conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
    r"""Applies a 1D convolution over an input signal composed of several input planes.

    See :class:`~tensorplay.nn.Conv1d` for details and output shape.

    Args:
        input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iW)`
        weight: filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kW)`
        bias: optional bias of shape :math:`(\text{out\_channels})`. Default: ``None``
        stride: the stride of the convolving kernel. Can be a single number or
          a one-element tuple `(sW,)`. Default: 1
        padding: implicit paddings on both sides of the input. Can be a single number, a one-element tuple `(padW,)`, ``'valid'`` for no padding, or ``'same'`` to keep every output position (unit stride only). Default: 0
        dilation: the spacing between kernel elements. Can be a single number or
          a one-element tuple `(dW,)`. Default: 1
        groups: split input into groups, :math:`\text{in\_channels}` should be divisible by
          the number of groups. Default: 1

    Examples::

        >>> inputs = tp.randn(33, 16, 30)
        >>> filters = tp.randn(20, 16, 5)
        >>> F.conv1d(inputs, filters)
    """
    if isinstance(padding, str):
        stride = _single(stride)
        dilation = _single(dilation)
        input, padding = _conv_padding_mode(input, padding, stride, dilation, weight, 1)
    captured = _capture_call(conv1d, (input, weight, bias, stride, padding, dilation, groups), {})
    if captured is not None:
        return captured
    stride = _single(stride)
    padding = _single(padding)
    dilation = _single(dilation)
    if bias is None:
        bias = Tensor()
    return _C.conv1d(input, weight, bias, stride, padding, dilation, groups)

[docs]
def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
    r"""Applies a 2D convolution over an input image composed of several input planes.

    See :class:`~tensorplay.nn.Conv2d` for details and output shape.

    Args:
        input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`
        weight: filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kH , kW)`
        bias: optional bias tensor of shape :math:`(\text{out\_channels})`. Default: ``None``
        stride: the stride of the convolving kernel. Can be a single number or a
          tuple `(sH, sW)`. Default: 1
        padding: implicit paddings on both sides of the input. Can be a single number, a tuple `(padH, padW)`, ``'valid'`` for no padding, or ``'same'`` to keep every output position (unit stride only). Default: 0
        dilation: the spacing between kernel elements. Can be a single number or
          a tuple `(dH, dW)`. Default: 1
        groups: split input into groups, both :math:`\text{in\_channels}` and :math:`\text{out\_channels}`
          should be divisible by the number of groups. Default: 1

    Examples::

        >>> # With square kernels and equal stride
        >>> filters = tp.randn(8, 4, 3, 3)
        >>> inputs = tp.randn(1, 4, 5, 5)
        >>> F.conv2d(inputs, filters, padding=1)
    """
    stride = _pair(stride)
    dilation = _pair(dilation)
    if isinstance(padding, str):
        input, padding = _conv_padding_mode(input, padding, stride, dilation, weight, 2)
    else:
        padding = _pair(padding)
    captured = _capture_call(
        conv2d,
        (input, weight, bias, stride, padding, dilation, groups),
        {},
    )
    if captured is not None:
        return captured
    if bias is None:
        bias = Tensor()
    return _C.conv2d(input, weight, bias, stride, padding, dilation, groups)

[docs]
def conv3d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
    r"""Applies a 3D convolution over an input image composed of several input planes.

    See :class:`~tensorplay.nn.Conv3d` for details and output shape.

    Args:
        input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iD, iH , iW)`
        weight: filters of shape :math:`(\text{out\_channels} , \frac{\text{in\_channels}}{\text{groups}} , kD, kH , kW)`
        bias: optional bias tensor of shape :math:`(\text{out\_channels})`. Default: ``None``
        stride: the stride of the convolving kernel. Can be a single number or a
          tuple `(sD, sH, sW)`. Default: 1
        padding: implicit paddings on both sides of the input. Can be a single number, a tuple `(padD, padH, padW)`, ``'valid'`` for no padding, or ``'same'`` to keep every output position (unit stride only). Default: 0
        dilation: the spacing between kernel elements. Can be a single number or
          a tuple `(dD, dH, dW)`. Default: 1
        groups: split input into groups, both :math:`\text{in\_channels}` and :math:`\text{out\_channels}`
          should be divisible by the number of groups. Default: 1

    Examples::

        >>> # With square kernels and equal stride
        >>> filters = tp.randn(8, 4, 3, 3, 3)
        >>> inputs = tp.randn(1, 4, 5, 5, 5)
        >>> F.conv3d(inputs, filters, padding=1)
    """
    if isinstance(padding, str):
        stride = _triple(stride)
        dilation = _triple(dilation)
        input, padding = _conv_padding_mode(input, padding, stride, dilation, weight, 3)
    captured = _capture_call(conv3d, (input, weight, bias, stride, padding, dilation, groups), {})
    if captured is not None:
        return captured
    stride = _triple(stride)
    padding = _triple(padding)
    dilation = _triple(dilation)
    if bias is None:
        bias = Tensor()
    return _C.conv3d(input, weight, bias, stride, padding, dilation, groups)

[docs]
def conv_transpose2d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1):
    captured = _capture_call(conv_transpose2d, (input, weight, bias, stride, padding, output_padding, groups, dilation), {})
    if captured is not None:
        return captured
    stride = _pair(stride)
    padding = _pair(padding)
    output_padding = _pair(output_padding)
    dilation = _pair(dilation)
    if bias is None:
        bias = Tensor()
    return _C.conv_transpose2d(input, weight, bias, stride, padding, output_padding, groups, dilation)

[docs]
def conv_transpose3d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1):
    stride = _triple(stride)
    padding = _triple(padding)
    output_padding = _triple(output_padding)
    dilation = _triple(dilation)
    if bias is None:
        bias = Tensor()
    return _C.conv_transpose3d(input, weight, bias, stride, padding, output_padding, groups, dilation)

[docs]
def conv_transpose1d(input, weight, bias=None, stride=1, padding=0, output_padding=0, groups=1, dilation=1):
    stride = _single(stride)
    padding = _single(padding)
    output_padding = _single(output_padding)
    dilation = _single(dilation)
    if bias is None:
        bias = Tensor()
    return _C.conv_transpose1d(input, weight, bias, stride, padding, output_padding, groups, dilation)

[docs]
def unfold(input, kernel_size, dilation=1, padding=0, stride=1):
    r"""
    """
    if input.dim() not in (3, 4):
        raise ValueError(
            f"unfold: expected 3D (unbatched) or 4D input, got {input.dim()}D")
    return _C.im2col(input, _pair(kernel_size), _pair(dilation), _pair(padding), _pair(stride))

[docs]
def fold(input, output_size, kernel_size, dilation=1, padding=0, stride=1):
    r"""Combine an array of sliding local blocks into a tensor containing
    """
    if input.dim() not in (2, 3):
        raise ValueError(
            f"fold: expected 2D (unbatched) or 3D input, got {input.dim()}D")
    return _C.col2im(input, _pair(output_size), _pair(kernel_size), _pair(dilation),
                     _pair(padding), _pair(stride))

[docs]
def conv_tbc(input, weight, bias=None, pad=0):
    r"""Applies a 1D convolution over an input of shape (T, B, C) along the
    ``(kernel_width, in_channels, out_channels)``; the math is a standard
    cross-channel conv1d after permuting to (B, C, T).
    """
    if input.dim() != 3:
        raise ValueError("conv_tbc: input must have 3 dimensions (T, B, C)")
    if weight.dim() != 3:
        raise ValueError(
            "Weight tensor must have 3 dims: kernel_width, in_channels, out_channels.")
    if weight.size(1) != input.size(2):
        raise ValueError(
            f"Input dim 2 (input channels) is not == dim 1 in the weight tensor")
    x = input.permute(1, 2, 0)              # (B, C_in, T)
    w = weight.permute(2, 1, 0)             # (C_out, C_in, k)
    if bias is None:
        bias = Tensor()
    out = _C.conv1d(x, w.contiguous(), bias, (1,), (pad,), (1,), 1)
    return out.permute(2, 0, 1)             # (T, B, C_out)

[docs]
def max_pool2d(input, kernel_size, stride=None, padding=0, dilation=1, ceil_mode=False, return_indices=False):
    captured = _capture_call(
        max_pool2d,
        (input, kernel_size, stride, padding, dilation, ceil_mode, return_indices),
        {},
    )
    if captured is not None:
        return captured
    if return_indices:
        return max_pool2d_with_indices(
            input, kernel_size, stride=stride, padding=padding,
            dilation=dilation, ceil_mode=ceil_mode)
    kernel_size = _pair(kernel_size)
    if stride is None:
        stride = kernel_size
    else:
        stride = _pair(stride)
    padding = _pair(padding)
    dilation = _pair(dilation)
    # Non-contiguous inputs (any memory format) are normalized by the CPU
    # kernel itself; channels-last tensors keep their layout end to end.
    return _C.max_pool2d(input, kernel_size, stride, padding,
                         dilation, ceil_mode)

[docs]
def avg_pool2d(input, kernel_size, stride=None, padding=0, ceil_mode=False, count_include_pad=True, divisor_override=None):
    captured = _capture_call(avg_pool2d, (input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override), {})
    if captured is not None:
        return captured
    kernel_size = _pair(kernel_size)
    if stride is None:
        stride = kernel_size
    else:
        stride = _pair(stride)
    padding = _pair(padding)
    return _C.avg_pool2d(input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override)

[docs]
def adaptive_avg_pool2d(input, output_size):
    # output_size can be int or (int, int) or (None, int) etc.
    output_size = _pair(output_size)
    captured = _capture_call(adaptive_avg_pool2d, (input, output_size), {})
    if captured is not None:
        return captured
    return _C.adaptive_avg_pool2d(input, output_size)

[docs]
def adaptive_max_pool2d(input, output_size):
    captured = _capture_call(adaptive_max_pool2d, (input, output_size), {})
    if captured is not None:
        return captured
    output_size = list(_pair(output_size))
    # Route through the (values, indices) op so autograd saves indices and the
    return _C.adaptive_max_pool2d_with_indices(input, output_size)[0]

# Normalization functions

[docs]
def batch_norm(input, running_mean=None, running_var=None, weight=None, bias=None, training=False, momentum=0.1, eps=1e-5):
    captured = _capture_call(
        batch_norm,
        (input, running_mean, running_var, weight, bias, training, momentum, eps),
        {},
    )
    if captured is not None:
        return captured
    return _C.batch_norm(input, weight, bias, running_mean, running_var, training, momentum, eps)

[docs]
def layer_norm(input, normalized_shape, weight=None, bias=None, eps=1e-5):
    captured = _capture_call(layer_norm, (input, normalized_shape, weight, bias, eps), {})
    if captured is not None:
        return captured
    normalized_shape = _single(normalized_shape)
    return _C.native_layer_norm(input, normalized_shape, weight, bias, eps)[0]

[docs]
def group_norm(input, num_groups, weight=None, bias=None, eps=1e-5):
    captured = _capture_call(group_norm, (input, num_groups, weight, bias, eps), {})
    if captured is not None:
        return captured
    if input.dim() < 2 or input.device.type != "cuda":
        return _C.group_norm(input, num_groups, weight, bias, eps)
    spatial = 1
    for dim in range(2, input.dim()):
        spatial *= input.size(dim)
    output, _, _ = _C.native_group_norm(
        input, weight, bias, input.size(0), input.size(1), spatial,
        num_groups, eps)
    return output

[docs]
def instance_norm(input, running_mean=None, running_var=None, weight=None, bias=None, use_input_stats=True, momentum=0.1, eps=1e-5):
    captured = _capture_call(instance_norm, (input, running_mean, running_var, weight, bias, use_input_stats, momentum, eps), {})
    if captured is not None:
        return captured
    return _C.instance_norm(input, weight, bias, running_mean, running_var, use_input_stats, momentum, eps)

[docs]
def pad(input, pad, mode='constant', value=0):
    r"""Pads tensor.  ``pad`` values are described starting from the last
    dimension and may cover any suffix of the input dimensions.
    """
    captured = _capture_call(globals()["pad"], (input, pad, mode, value), {})
    if captured is not None:
        return captured
    if mode == 'constant':
        return _C.constant_pad_nd(input, list(pad), value)
    ndim = input.dim()
    pad = list(pad)
    if len(pad) % 2 != 0:
        raise ValueError(f"padding length must be even, got {len(pad)}")
    if len(pad) > 2 * ndim:
        raise ValueError(
            f"padding length {len(pad)} exceeds the {ndim}-D input"
        )
    if mode == 'reflect':
        return _C.reflection_pad_nd(input, pad)
    if mode == 'replicate':
        return _C.replication_pad_nd(input, pad)
    if mode == 'circular':
        return _C.circular_pad_nd(input, pad)
    raise ValueError(f"Padding mode '{mode}' not supported")

# Loss functions

[docs]
def mse_loss(input, target, reduction='mean'):
    captured = _capture_call(mse_loss, (input, target, reduction), {})
    if captured is not None:
        return captured
    if not (target.size() == input.size()):
        print(f"Warning: Using a target size ({target.size()}) that is different to the input size ({input.size()}). "
              "This will likely lead to incorrect results due to broadcasting. "
              "Please ensure they have the same size.")
    
    reduction_enum = 1
    if reduction == 'none': reduction_enum = 0
    elif reduction == 'mean': reduction_enum = 1
    elif reduction == 'sum': reduction_enum = 2
    else: raise ValueError(f"{reduction} is not a valid value for reduction")

    return _C.mse_loss(input, target, reduction_enum)

def _nll_loss_red_enum(reduction):
    if reduction == 'none': return 0
    elif reduction == 'mean': return 1
    elif reduction == 'sum': return 2
    raise ValueError(f"{reduction} is not a valid value for reduction")

[docs]
def nll_loss(input, target, weight=None, size_average=None, ignore_index=-100,
             reduce=None, reduction='mean'):
    r"""The negative log likelihood loss.

    ``input`` with a scalar ``target``, 2D ``(N, C)``, and N-d
    ``target``.

    See :class:`~tensorplay.nn.NLLLoss` for details.
    """
    captured = _capture_call(nll_loss, (input, target, weight, size_average, ignore_index, reduce, reduction), {})
    if captured is not None:
        return captured
    if size_average is not None or reduce is not None:
        if size_average is None: size_average = True
        if reduce is None: reduce = True
        if not reduce: reduction = 'none'
        elif size_average: reduction = 'mean'
        else: reduction = 'sum'

    reduction_enum = _nll_loss_red_enum(reduction)

    if input.dim() <= 2:
        # nll_loss returns (output, total_weight)
        output, _ = _C.nll_loss(input, target, weight, reduction_enum, ignore_index)
        return output

    if input.dim() == 4:
        # input with (N, H, W) target; autograd flows through
        # nll_loss2d_backward.
        t = target if target.dtype == DType.int64 else target.to(DType.int64)
        output, _ = _C.nll_loss2d(input, t, weight, reduction_enum, ignore_index)
        return output

    # every spatial position acts as its own batch row, so move classes last,
    # flatten to (-1, C), run the 2-D kernel and restore the target shape.
    if tuple(target.size())[1:] != tuple(input.size())[2:]:
        expected = tuple(input.size()[:1] + input.size()[2:])
        raise ValueError(f"Expected target size {expected}, got {tuple(target.size())}")
    n = input.size(0)
    c = input.size(1)
    # (N, C, d_1, ..., d_k) -> (N, d_1, ..., d_k, C) -> (N * prod(d_i), C);
    # row order matches target's contiguous flattening.
    x = input.permute([0] + list(range(2, input.dim())) + [1]).contiguous().reshape(-1, c)
    t = target.contiguous().reshape(-1)
    if t.dtype != DType.int64:
        t = t.to(DType.int64)
    output, _ = _C.nll_loss(x, t, weight, reduction_enum, ignore_index)
    if reduction == 'none':
        output = output.reshape(target.size())
    return output

[docs]
def cross_entropy(input, target, weight=None, size_average=None, ignore_index=-100,
                  reduce=None, reduction='mean', label_smoothing=0.0):
    r"""Compute the cross entropy loss between input logits and target.

    the class-probability path, positive ``label_smoothing`` blends the NLL
    with a smoothed uniform term, and otherwise this is
    ``nll_loss(log_softmax(input), target)`` with N-d support.

    See :class:`~tensorplay.nn.CrossEntropyLoss` for details.
    """
    captured = _capture_call(cross_entropy, (input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing), {})
    if captured is not None:
        return captured
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)

    # Hot path for the common training-loop case ((N, C) logits, int64
    # class targets, no class weights / label smoothing): go straight to
    # the native ops, skipping the legacy-reduction and N-d shape handling.
    if (weight is None and label_smoothing == 0.0 and input.dim() == 2
            and target.dtype == tensorplay.int64):
        red = 1 if reduction == 'mean' else (
            2 if reduction == 'sum' else (0 if reduction == 'none' else -1))
        if red >= 0:
            output, _ = _C.nll_loss(
                _C.log_softmax(input, 1, tensorplay.undefined),
                target, None, red, ignore_index)
            return output

    class_dim = 0 if input.dim() == 1 else 1
    n_classes = input.size(class_dim)
    if weight is not None and (weight.dim() != 1 or weight.numel() != n_classes):
        raise ValueError(
            f"cross_entropy: weight tensor should be defined either for all "
            f"{n_classes} classes or no classes but got weight tensor of "
            f"shape: {tuple(weight.size())}")

    if tuple(target.size()) == tuple(input.size()):
        # Soft targets when input and target shapes are the same
        # Handle matching input and target shapes as class probabilities.
        if ignore_index >= 0:
            raise ValueError("ignore_index is not supported for floating point target")
        if label_smoothing > 1.0:
            raise ValueError(f"label_smoothing must be between 0.0 and 1.0. Got: {label_smoothing}")
        input_ = log_softmax(input, class_dim)
        if label_smoothing > 0.0:
            target = target * (1 - label_smoothing) + label_smoothing / n_classes
        if weight is not None:
            w_shape = [1] * input.dim()
            w_shape[class_dim] = n_classes
            loss = -(input_ * target * weight.view(w_shape)).sum(class_dim)
        else:
            loss = -(input_ * target).sum(class_dim)
        if reduction == "none":
            return loss
        total = loss.sum()
        if reduction == "sum":
            return total
        if input.numel() == 0:
            return tensorplay.full([], float("nan"), dtype=total.dtype, device=total.device)
        return total / (input.numel() / n_classes)

    if label_smoothing > 0.0:
        # Blend the class loss with a uniform class distribution.
        if label_smoothing > 1.0:
            raise ValueError(f"label_smoothing must be between 0.0 and 1.0. Got: {label_smoothing}")
        input_ = log_softmax(input, class_dim)
        nllloss = nll_loss(input_, target, weight, None, ignore_index, None, reduction)
        if weight is not None:
            w_shape = [1] * input_.dim()
            w_shape[class_dim] = n_classes
            smooth_loss = -(input_ * weight.view(w_shape)).sum(class_dim)
        else:
            smooth_loss = -input_.sum(class_dim)
        ignore_mask = target.eq(ignore_index)
        smooth_loss = smooth_loss.masked_fill(ignore_mask, 0.0)
        if reduction == "mean":
            if weight is not None:
                filtered_target = target.masked_fill(ignore_mask, 0)
                tgt_weights = weight.index_select(0, filtered_target.reshape(-1))
                weight_sum = tgt_weights.masked_fill_(ignore_mask.reshape(-1), 0).sum()
                ret = smooth_loss.sum() / weight_sum
            else:
                true_mask = tensorplay.logical_not(ignore_mask)
                ret = smooth_loss.sum() / true_mask.to(smooth_loss.dtype).sum()
        elif reduction == "sum":
            ret = smooth_loss.sum()
        elif reduction == "none":
            ret = smooth_loss
        else:
            raise ValueError(f"{reduction} is not valid")
        return (1 - label_smoothing) * nllloss + ret * (label_smoothing / n_classes)

    return nll_loss(log_softmax(input, class_dim), target, weight, None, ignore_index, None, reduction)

# -----------------------------------------------------------------------------
# Activation / misc functions.  These are thin wrappers over the native
# dispatcher ops declared in the native schema; the element-wise
# -----------------------------------------------------------------------------

[docs]
def gelu(input: Tensor, approximate: str = 'none') -> Tensor:
    r"""gelu(input, approximate='none') -> Tensor

    When `approximate` is 'none', applies
    :math:`\text{GELU}(x) = x * \Phi(x)`; 'tanh' uses the tanh estimation.
    """
    captured = _capture_call(gelu, (input,), {"approximate": approximate})
    if captured is not None:
        return captured
    return tensorplay._C.gelu(self=input, approximate=approximate)

[docs]
def relu6(input: Tensor, inplace: bool = False) -> Tensor:
    r"""relu6(input, inplace=False) -> Tensor

    """
    out = tensorplay.relu6(input)
    return input.copy_(out) if inplace else out

[docs]
def hardswish(input: Tensor, inplace: bool = False) -> Tensor:
    r"""hardswish(input, inplace=False) -> Tensor"""
    out = tensorplay.hardswish(input)
    return input.copy_(out) if inplace else out

[docs]
def hardsigmoid(input: Tensor, inplace: bool = False) -> Tensor:
    r"""hardsigmoid(input, inplace=False) -> Tensor"""
    out = tensorplay.hardsigmoid(input)
    return input.copy_(out) if inplace else out

[docs]
def leaky_relu(input: Tensor, negative_slope: float = 0.01, inplace: bool = False) -> Tensor:
    r"""leaky_relu(input, negative_slope=0.01, inplace=False) -> Tensor"""
    out = tensorplay.leaky_relu(input, negative_slope)
    return input.copy_(out) if inplace else out

[docs]
def softplus(input: Tensor, beta: float = 1.0, threshold: float = 20.0) -> Tensor:
    r"""softplus(input, beta=1, threshold=20) -> Tensor"""
    return tensorplay.softplus(input, beta, threshold)

[docs]
def elu(input: Tensor, alpha: float = 1.0, inplace: bool = False) -> Tensor:
    r"""elu(input, alpha=1, inplace=False) -> Tensor"""
    out = tensorplay.elu(input, alpha)
    return input.copy_(out) if inplace else out

[docs]
def mish(input: Tensor, inplace: bool = False) -> Tensor:
    r"""mish(input, inplace=False) -> Tensor"""
    out = tensorplay.mish(input)
    return input.copy_(out) if inplace else out

[docs]
def selu(input: Tensor, inplace: bool = False) -> Tensor:
    r"""selu(input, inplace=False) -> Tensor"""
    out = tensorplay.selu(input)
    return input.copy_(out) if inplace else out

[docs]
def celu(input: Tensor, alpha: float = 1.0, inplace: bool = False) -> Tensor:
    r"""celu(input, alpha=1, inplace=False) -> Tensor"""
    out = tensorplay.celu(input, alpha)
    return input.copy_(out) if inplace else out

threshold_ = _add_docstr(
    _C.threshold_,
    r"""
threshold_(input, threshold, value) -> Tensor

In-place version of :func:`~threshold`.
""",
)

relu_ = _add_docstr(
    _C.relu_,
    r"""
relu_(input) -> Tensor

In-place version of :func:`~relu`.
""",
)

hardtanh_ = _add_docstr(
    _C.hardtanh_,
    r"""
hardtanh_(input, min_val=-1., max_val=1.) -> Tensor

In-place version of :func:`~hardtanh`.
""",
)

elu_ = _add_docstr(
    _C.elu_,
    r"""
elu_(input, alpha=1.) -> Tensor

In-place version of :func:`~elu`.
""",
)

selu_ = _add_docstr(
    _C.selu_,
    r"""
selu_(input) -> Tensor

In-place version of :func:`~selu`.
""",
)

celu_ = _add_docstr(
    _C.celu_,
    r"""
celu_(input, alpha=1.) -> Tensor

In-place version of :func:`~celu`.
""",
)

leaky_relu_ = _add_docstr(
    _C.leaky_relu_,
    r"""
leaky_relu_(input, negative_slope=0.01) -> Tensor

In-place version of :func:`~leaky_relu`.
""",
)

[docs]
def lstm_cell(input, hx, cx, w_ih, w_hh, b_ih=None, b_hh=None):
    r"""lstm_cell(input, hx, cx, w_ih, w_hh, b_ih=None, b_hh=None) -> (Tensor, Tensor)

    One time step of a long short-term memory cell. Returns the next
    hidden state and next cell state.
    """
    return _C.lstm_cell(input, hx, cx, w_ih, w_hh, b_ih, b_hh)

[docs]
def rnn_relu_cell(input, hx, w_ih, w_hh, b_ih=None, b_hh=None):
    r"""rnn_relu_cell(input, hx, w_ih, w_hh, b_ih=None, b_hh=None) -> Tensor

    One time step of an Elman RNN cell with ReLU nonlinearity.
    """
    return _C.rnn_relu_cell(input, hx, w_ih, w_hh, b_ih, b_hh)

[docs]
def rnn_tanh_cell(input, hx, w_ih, w_hh, b_ih=None, b_hh=None):
    r"""rnn_tanh_cell(input, hx, w_ih, w_hh, b_ih=None, b_hh=None) -> Tensor

    One time step of an Elman RNN cell with tanh nonlinearity.
    """
    return _C.rnn_tanh_cell(input, hx, w_ih, w_hh, b_ih, b_hh)

[docs]
def gru_cell(input, hx, w_ih, w_hh, b_ih=None, b_hh=None):
    r"""gru_cell(input, hx, w_ih, w_hh, b_ih=None, b_hh=None) -> Tensor

    One time step of a gated recurrent unit cell.
    """
    return _C.gru_cell(input, hx, w_ih, w_hh, b_ih, b_hh)

[docs]
def glu(input: Tensor, dim: int = -1) -> Tensor:
    r"""glu(input, dim=-1) -> Tensor

    Gated Linear Unit: :math:`a * \sigma(b)` where the input is split in half
    along ``dim``.
    """
    return tensorplay.glu(input, dim)

[docs]
def normalize(input: Tensor, p: float = 2.0, dim: int = 1, eps: float = 1e-12) -> Tensor:
    r"""normalize(input, p=2, dim=1, eps=1e-12) -> Tensor

    Performs :math:`L_p` normalization over the specified dimension —
    """
    denom = input.norm([dim], p, True).clamp_min(eps)
    return input / denom

def one_hot(tensor: Tensor, num_classes: int = -1) -> Tensor:
    r"""one_hot(tensor, num_classes=-1) -> Tensor"""
    return tensorplay.one_hot(tensor, num_classes)

[docs]
def interpolate(
    input: Tensor,
    size=None,
    scale_factor=None,
    mode: str = 'nearest',
    align_corners=None,
    recompute_scale_factor=None,
    antialias: bool = False,
) -> Tensor:
    r"""interpolate(input, size=None, scale_factor=None, mode='nearest',
    align_corners=None) -> Tensor

    Routes to the native ``upsample_*`` ops exactly like
    """
    captured = _capture_call(interpolate, (input, size, scale_factor, mode, align_corners, recompute_scale_factor, antialias), {})
    if captured is not None:
        return captured
    if input.is_quantized():
        from tensorplay.ao.nn.quantized.functional import interpolate as quantized_interpolate
        return quantized_interpolate(input, size, scale_factor, mode, align_corners,
                                    recompute_scale_factor, antialias)
    if size is None and scale_factor is None:
        raise ValueError("need to define size or scale_factor")
    if size is not None and scale_factor is not None:
        raise ValueError("only one of size or scale_factor should be defined")

    ndim = input.dim()
    if ndim < 3:
        raise ValueError(f"interpolate expects at least 3D input, got {ndim}D")
    spatial = ndim - 2
    align_corners_provided = align_corners is not None
    # mode branches can operate on a concrete size.
    if size is None and scale_factor is not None:
        if isinstance(scale_factor, (int, float)):
            scale_list = [scale_factor] * spatial
        else:
            scale_list = list(scale_factor)
            if len(scale_list) != spatial:
                raise ValueError(
                    f"scale_factor must have {spatial} values, got {len(scale_list)}"
                )
        import math
        if any(float(s) <= 0 for s in scale_list):
            raise ValueError("scale_factor values must be positive")
        size = [int(math.floor(float(input.size(2 + i)) * s))
                for i, s in enumerate(scale_list)]
    elif isinstance(size, int):
        size = [size] * spatial
    else:
        size = list(size)
        if len(size) != spatial:
            raise ValueError(
                f"size must have {spatial} values, got {len(size)}"
            )
    if any(int(s) <= 0 for s in size):
        raise ValueError(f"interpolate output sizes must be positive, got {size}")
    if mode in ('nearest', 'nearest-exact'):
        if align_corners is not None:
            raise ValueError("align_corners option can only be set with interpolating modes")
        # 'nearest-exact' resolves to the pixel-center kernels; the legacy
        # 'nearest' keeps the deprecated asymmetric-index behavior.
        if mode == 'nearest-exact':
            import tensorplay.functional as _functional
            if ndim == 3:
                return _functional._upsample_nearest_exact1d(input, size, None)
            elif ndim == 4:
                return _functional._upsample_nearest_exact2d(input, size, None)
            elif ndim == 5:
                return _functional._upsample_nearest_exact3d(input, size, None)
            raise ValueError(f"Expected 3D, 4D or 5D input, got {ndim}D")
        if ndim == 3:
            return tensorplay.upsample_nearest1d(input, size)
        elif ndim == 4:
            return tensorplay.upsample_nearest2d(input, size)
        elif ndim == 5:
            return tensorplay.upsample_nearest3d(input, size)
        raise ValueError(f"Expected 3D, 4D or 5D input, got {ndim}D")

    if align_corners is None:
        align_corners = False

    import math
    if scale_factor is not None:
        if isinstance(scale_factor, (int, float)):
            scale_factor = [scale_factor] * spatial
        elif len(scale_factor) != spatial:
            raise ValueError(
                f"scale_factor must have {spatial} values, got {len(scale_factor)}"
            )
        if any(float(f) <= 0 for f in scale_factor):
            raise ValueError("scale_factor values must be positive")
        if recompute_scale_factor:
            size = [int(math.floor(float(input.shape[2 + i]) * f)) for i, f in enumerate(scale_factor)]

    if mode == 'area':
        if align_corners_provided:
            raise ValueError("align_corners option cannot be set for area interpolation")
        if ndim == 3:
            return adaptive_avg_pool1d(input, size[0])
        if ndim == 4:
            return adaptive_avg_pool2d(input, size)
        if ndim == 5:
            return adaptive_avg_pool3d(input, size)
        raise ValueError(f"Expected 3D, 4D or 5D input, got {ndim}D")

    if mode == 'linear':
        if ndim != 3:
            raise ValueError("linear interpolation expects 3D input")
        if antialias:
            raise NotImplementedError("interpolate: antialias is not supported with mode='linear'")
        return tensorplay.upsample_linear1d(input, size, align_corners)
    elif mode == 'bilinear':
        if ndim != 4:
            raise ValueError("bilinear interpolation expects 4D input")
        if antialias:
            return _C._upsample_bilinear2d_aa(input, size, align_corners, None)
        return tensorplay.upsample_bilinear2d(input, size, align_corners)
    elif mode == 'bicubic':
        if ndim != 4:
            raise ValueError("bicubic interpolation expects 4D input")
        if antialias:
            return _C._upsample_bicubic2d_aa(input, size, align_corners, None)
        return tensorplay.upsample_bicubic2d(input, size, align_corners)
    elif mode == 'trilinear':
        if ndim != 5:
            raise ValueError("trilinear interpolation expects 5D input")
        if antialias:
            raise NotImplementedError("interpolate: antialias is not supported with mode='trilinear'")
        return tensorplay.upsample_trilinear3d(input, size, align_corners)
    raise ValueError(f"interpolate: mode '{mode}' is not supported")

[docs]
def multi_head_attention_forward(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    embed_dim_to_check: int,
    num_heads: int,
    in_proj_weight: Tensor = None,
    in_proj_bias: Tensor = None,
    bias_k=None,
    bias_v=None,
    add_zero_attn: bool = False,
    dropout_p: float = 0.0,
    out_proj_weight: Tensor = None,
    out_proj_bias: Tensor = None,
    training: bool = True,
    key_padding_mask=None,
    need_weights: bool = True,
    attn_mask=None,
    use_separate_proj_weight: bool = False,
    q_proj_weight=None,
    k_proj_weight=None,
    v_proj_weight=None,
    static_k=None,
    static_v=None,
    average_attn_weights: bool = True,
    is_causal: bool = False,
):
    r"""multi_head_attention_forward(query, key, value, embed_dim_to_check,
    num_heads, in_proj_weight, in_proj_bias=None, bias_k=None, bias_v=None,
    add_zero_attn=False, dropout_p=0.0, out_proj_weight=None,
    out_proj_bias=None, training=True, key_padding_mask=None,
    need_weights=True, attn_mask=None, use_separate_proj_weight=False,
    q_proj_weight=None, k_proj_weight=None, v_proj_weight=None,
    static_k=None, static_v=None, average_attn_weights=True,
    is_causal=False) -> (Tensor, Optional[Tensor])

    Computes multi-head attention on (L, N, E) inputs (2D unbatched inputs are
    promoted internally and the batch dim is squeezed on return). Returns the
    projected output of shape (L, N, E) and, when ``need_weights`` is true,
    the attention weights of shape (N, L, S) — or (num_heads, L, S) with
    ``average_attn_weights=False``.
    """
    # Unbatched inputs carry a temporary batch dim; outputs squeeze it back.
    if query.dim() == 3:
        is_batched = True
        if key.dim() != 3 or value.dim() != 3:
            raise AssertionError(
                "for batched query, key and value must also be 3D")
        if key_padding_mask is not None and key_padding_mask.dim() > 2:
            raise AssertionError(
                "key_padding_mask must be 1D or 2D for batched input")
        if attn_mask is not None and attn_mask.dim() not in (2, 3):
            raise AssertionError(
                f"attn_mask must be 2D or 3D for batched input, got {attn_mask.dim()}D")
    elif query.dim() == 2:
        is_batched = False
        if key.dim() != 2 or value.dim() != 2:
            raise AssertionError(
                "for unbatched query, key and value must also be 2D")
        if key_padding_mask is not None and key_padding_mask.dim() != 1:
            raise AssertionError(
                "key_padding_mask must be 1D for unbatched input")
        if attn_mask is not None and attn_mask.dim() != 2:
            raise AssertionError(
                f"attn_mask must be 2D for unbatched input, got {attn_mask.dim()}D")
        query = query.unsqueeze(1)
        key = key.unsqueeze(1)
        value = value.unsqueeze(1)
        if key_padding_mask is not None:
            key_padding_mask = key_padding_mask.unsqueeze(0)
    else:
        raise AssertionError(
            f"query has to be 2d or 3d, but got {query.dim()}d")

    tgt_len, bsz, embed_dim = query.shape
    if num_heads <= 0:
        raise AssertionError(f"num_heads must be positive, got {num_heads}")
    if embed_dim_to_check != embed_dim:
        raise AssertionError(
            f"was expecting embedding dimension of {embed_dim_to_check}, "
            f"but got {embed_dim}")
    if key.shape[1] != bsz or value.shape[1] != bsz:
        raise AssertionError(
            "key and value batch dimensions must match the query batch "
            f"dimension ({bsz})")
    src_len = key.shape[0]
    head_dim = embed_dim // num_heads
    if head_dim * num_heads != embed_dim:
        raise AssertionError(
            f"embed_dim {embed_dim} not divisible by num_heads {num_heads}")

    if use_separate_proj_weight:
        if key.shape[:2] != value.shape[:2]:
            raise AssertionError(
                f"key's sequence and batch dims {tuple(key.shape[:2])} do not "
                f"match value's {tuple(value.shape[:2])}")
    elif key.shape != value.shape:
        raise AssertionError(
            f"key shape {tuple(key.shape)} does not match value shape {tuple(value.shape)}")

    # Bool masks become additive float masks so they can merge by addition.
    key_padding_mask = _canonical_mask(
        mask=key_padding_mask,
        mask_name="key_padding_mask",
        other_type=_none_or_dtype(attn_mask),
        other_name="attn_mask",
        target_type=query.dtype,
    )

    if is_causal and attn_mask is None and need_weights:
        causal_rows = tensorplay.arange(
            tgt_len, dtype=DType.int64, device=query.device).view(tgt_len, 1)
        causal_columns = tensorplay.arange(
            src_len, dtype=DType.int64, device=query.device).view(1, src_len)
        causal_fill = tensorplay.full(
            [tgt_len, src_len], float("-inf"), dtype=query.dtype,
            device=query.device)
        causal_mask = causal_rows < causal_columns
        attn_mask = tensorplay.where(
            causal_mask, causal_fill,
            tensorplay.zeros([tgt_len, src_len], dtype=query.dtype,
                             device=query.device))
        is_causal = False
    elif is_causal and attn_mask is None:
        raise RuntimeError(
            "Need attn_mask if specifying the is_causal hint. "
            "You may use the Transformer module method "
            "`generate_square_subsequent_mask` to create this mask.")

    if is_causal and key_padding_mask is None and not need_weights:
        # No mask fusion needed: pass the is_causal hint straight to SDPA.
        attn_mask = None
    else:
        attn_mask = _canonical_mask(
            mask=attn_mask,
            mask_name="attn_mask",
            other_type=None,
            other_name="",
            target_type=query.dtype,
            check_other=False,
        )
        if key_padding_mask is not None:
            # The merged mask is no longer causal.
            is_causal = False

    # compute in-projection
    if not use_separate_proj_weight:
        if in_proj_weight is None:
            raise AssertionError(
                "use_separate_proj_weight is False but in_proj_weight is None")
        q, k, v = _in_projection_packed(query, key, value, in_proj_weight, in_proj_bias)
    else:
        if q_proj_weight is None:
            raise AssertionError(
                "use_separate_proj_weight is True but q_proj_weight is None")
        if k_proj_weight is None:
            raise AssertionError(
                "use_separate_proj_weight is True but k_proj_weight is None")
        if v_proj_weight is None:
            raise AssertionError(
                "use_separate_proj_weight is True but v_proj_weight is None")
        if in_proj_bias is None:
            b_q = b_k = b_v = None
        else:
            b_q, b_k, b_v = tensorplay.chunk(in_proj_bias, 3)
        q, k, v = _in_projection(
            query, key, value, q_proj_weight, k_proj_weight, v_proj_weight,
            b_q, b_k, b_v,
        )

    # prep attention mask: promote a 2D mask to 3D (broadcast over batch).
    if attn_mask is not None:
        if attn_mask.dim() == 2:
            correct_2d_size = (tgt_len, src_len)
            if tuple(attn_mask.shape) != correct_2d_size:
                raise RuntimeError(
                    f"The shape of the 2D attn_mask is {tuple(attn_mask.shape)}, "
                    f"but should be {correct_2d_size}.")
            attn_mask = attn_mask.unsqueeze(0)
        elif attn_mask.dim() == 3:
            correct_3d_size = (bsz * num_heads, tgt_len, src_len)
            if tuple(attn_mask.shape) != correct_3d_size:
                raise RuntimeError(
                    f"The shape of the 3D attn_mask is {tuple(attn_mask.shape)}, "
                    f"but should be {correct_3d_size}.")
        else:
            raise RuntimeError(
                f"attn_mask's dimension {attn_mask.dim()} is not supported")

    # extra bias heads join the key/value sequences (dim 0, pre-reshape).
    if bias_k is not None and bias_v is not None:
        if static_k is not None:
            raise AssertionError("bias cannot be added to static key.")
        if static_v is not None:
            raise AssertionError("bias cannot be added to static value.")
        k = tensorplay.cat([k, bias_k.repeat(1, bsz, 1)])
        v = tensorplay.cat([v, bias_v.repeat(1, bsz, 1)])
        if attn_mask is not None:
            attn_mask = pad(attn_mask, (0, 1))
        if key_padding_mask is not None:
            key_padding_mask = pad(key_padding_mask, (0, 1))
    else:
        if bias_k is not None:
            raise AssertionError("bias_k is set but bias_v is None")
        if bias_v is not None:
            raise AssertionError("bias_v is set but bias_k is None")

    # reshape to heads-first (num_heads folded into the batch dim)
    q = q.reshape(tgt_len, bsz * num_heads, head_dim).transpose(0, 1)
    if static_k is None:
        k = k.reshape(k.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
    else:
        if static_k.shape[0] != bsz * num_heads:
            raise AssertionError(
                f"expecting static_k.size(0) of {bsz * num_heads}, "
                f"but got {static_k.shape[0]}")
        if static_k.shape[2] != head_dim:
            raise AssertionError(
                f"expecting static_k.size(2) of {head_dim}, but got {static_k.shape[2]}")
        k = static_k
    if static_v is None:
        v = v.reshape(v.shape[0], bsz * num_heads, head_dim).transpose(0, 1)
    else:
        if static_v.shape[0] != bsz * num_heads:
            raise AssertionError(
                f"expecting static_v.size(0) of {bsz * num_heads}, "
                f"but got {static_v.shape[0]}")
        if static_v.shape[2] != head_dim:
            raise AssertionError(
                f"expecting static_v.size(2) of {head_dim}, but got {static_v.shape[2]}")
        v = static_v

    # extra zero key/value joined along the (heads-first) sequence dim.
    if add_zero_attn:
        zero_attn_shape = (bsz * num_heads, 1, head_dim)
        k = tensorplay.cat(
            [k, tensorplay.zeros(zero_attn_shape, dtype=k.dtype, device=k.device)],
            dim=1,
        )
        v = tensorplay.cat(
            [v, tensorplay.zeros(zero_attn_shape, dtype=v.dtype, device=v.device)],
            dim=1,
        )
        if attn_mask is not None:
            attn_mask = pad(attn_mask, (0, 1))
        if key_padding_mask is not None:
            key_padding_mask = pad(key_padding_mask, (0, 1))

    src_len = k.shape[1]

    # merge key padding into the attention mask (broadcast over query heads)
    if key_padding_mask is not None:
        if key_padding_mask.shape[0] != bsz:
            raise AssertionError(
                f"Expected key_padded_mask.shape[0] to be {bsz}, "
                f"but got {key_padding_mask.shape[0]}")
        if key_padding_mask.shape[1] != src_len:
            raise AssertionError(
                f"Expected key_padded_mask.shape[1] to be {src_len}, "
                f"but got {key_padding_mask.shape[1]}")
        key_padding_mask = (
            key_padding_mask.reshape(bsz, 1, 1, src_len)
            .expand(-1, num_heads, -1, -1)
            .reshape(bsz * num_heads, 1, src_len)
        )
        if attn_mask is None:
            attn_mask = key_padding_mask
        else:
            attn_mask = attn_mask + key_padding_mask

    if not training:
        dropout_p = 0.0

    if need_weights:
        E = q.shape[2]
        q_scaled = q * math.sqrt(1.0 / float(E))

        if attn_mask is not None:
            attn_output_weights = tensorplay.baddbmm(
                attn_mask.to(q_scaled.dtype), q_scaled, k.transpose(-2, -1)
            )
        else:
            attn_output_weights = tensorplay.bmm(q_scaled, k.transpose(-2, -1))
        attn_output_weights = softmax(attn_output_weights, dim=-1)
        if dropout_p > 0.0:
            attn_output_weights = dropout(attn_output_weights, p=dropout_p)

        attn_output = tensorplay.bmm(attn_output_weights, v)
        attn_output = attn_output.transpose(0, 1).reshape(tgt_len * bsz, embed_dim)
        attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
        attn_output = attn_output.reshape(tgt_len, bsz, attn_output.shape[1])

        attn_output_weights = attn_output_weights.reshape(bsz, num_heads, tgt_len, src_len)
        if average_attn_weights:
            attn_output_weights = attn_output_weights.mean(dim=1)

        if not is_batched:
            attn_output = attn_output.squeeze(1)
            attn_output_weights = attn_output_weights.squeeze(0)
        return attn_output, attn_output_weights
    else:
        # (1, L, S) masks broadcast per batch; (N*num_heads, L, S) fold to 4D.
        if attn_mask is not None:
            if attn_mask.shape[0] == 1 and attn_mask.dim() == 3:
                attn_mask = attn_mask.unsqueeze(0)
            else:
                attn_mask = attn_mask.reshape(bsz, num_heads, -1, src_len)

        q = q.reshape(bsz, num_heads, tgt_len, head_dim)
        k = k.reshape(bsz, num_heads, src_len, head_dim)
        v = v.reshape(bsz, num_heads, src_len, head_dim)

        attn_output = scaled_dot_product_attention(
            q, k, v, attn_mask, dropout_p, is_causal
        )
        attn_output = attn_output.permute(2, 0, 1, 3).reshape(bsz * tgt_len, embed_dim)

        attn_output = linear(attn_output, out_proj_weight, out_proj_bias)
        attn_output = attn_output.reshape(tgt_len, bsz, attn_output.shape[1])
        if not is_batched:
            attn_output = attn_output.squeeze(1)
        return attn_output, None

# -----------------------------------------------------------------------------
# dispatcher ops available here, following the formulas in
# -----------------------------------------------------------------------------

def _get_reduction_enum(reduction: str) -> int:
    if reduction == 'none':
        return 0
    elif reduction == 'mean':
        return 1
    elif reduction == 'sum':
        return 2
    raise ValueError(f"{reduction} is not valid")

def _legacy_get_string(size_average, reduce):
    if size_average is None:
        size_average = True
    if reduce is None:
        reduce = True
    if size_average and reduce:
        return 'mean'
    elif reduce:
        return 'sum'
    else:
        return 'none'

[docs]
def logsigmoid(input: Tensor) -> Tensor:
    r"""logsigmoid(input) -> Tensor

    Applies element-wise :math:`\text{LogSigmoid}(x_i) = \log \left(\frac{1}{1 + \exp(-x_i)}\right)`

    See :class:`~tensorplay.nn.LogSigmoid` for more details.
    """
    # autograd flows through log_sigmoid_backward.
    return _C.log_sigmoid(input)

[docs]
def softmin(input: Tensor, dim: Optional[int] = None, dtype=None) -> Tensor:
    r"""Apply a softmin function.

    Note that :math:`\text{Softmin}(x) = \text{Softmax}(-x)`.

    See :class:`~tensorplay.nn.Softmin` for more details.
    """
    if dim is None:
        dim = input.dim() - 1 if input.dim() in (0, 1, 3) else 1
    return softmax(-input, dim=dim, dtype=dtype)

[docs]
def softsign(input: Tensor) -> Tensor:
    r"""softsign(input) -> Tensor

    Applies element-wise, the function :math:`\text{SoftSign}(x) = \frac{x}{1 + |x|}`

    See :class:`~tensorplay.nn.Softsign` for more details.
    """
    return input / (input.abs() + 1)

[docs]
def tanhshrink(input: Tensor) -> Tensor:
    r"""tanhshrink(input) -> Tensor

    Applies element-wise, :math:`\text{Tanhshrink}(x) = x - \text{Tanh}(x)`

    See :class:`~tensorplay.nn.Tanhshrink` for more details.
    """
    return input - tensorplay.tanh(input)

[docs]
def hardtanh(
    input: Tensor, min_val: float = -1.0, max_val: float = 1.0,
    inplace: bool = False,
) -> Tensor:
    r"""hardtanh(input, min_val=-1.0, max_val=1.0, inplace=False) -> Tensor"""
    result = tensorplay.hardtanh(input, min_val=min_val, max_val=max_val)
    if inplace:
        return input.copy_(result)
    return result

[docs]
def hardshrink(input: Tensor, lambd: float = 0.5) -> Tensor:
    r"""hardshrink(input, lambd=0.5) -> Tensor

    Applies the hard shrinkage function element-wise.

    See :class:`~tensorplay.nn.Hardshrink` for more details.
    """
    return tensorplay.hardshrink(input, lambd)

[docs]
def softshrink(input: Tensor, lambd: float = 0.5) -> Tensor:
    r"""softshrink(input, lambd=0.5) -> Tensor

    Applies the soft shrinkage function element-wise.

    See :class:`~tensorplay.nn.Softshrink` for more details.
    """
    return tensorplay.softshrink(input, lambd)

[docs]
def rrelu(
    input: Tensor,
    lower: float = 1.0 / 8,
    upper: float = 1.0 / 3,
    training: bool = False,
    inplace: bool = False,
) -> Tensor:
    r"""rrelu(input, lower=1./8, upper=1./3, training=False, inplace=False) -> Tensor

    Randomized leaky ReLU.

    See :class:`~tensorplay.nn.RReLU` for more details.
    """
    if inplace:
        return _C.rrelu_(input, lower, upper, training)
    return _C.rrelu(input, lower, upper, training)

# -----------------------------------------------------------------------------
# Pooling helpers
# -----------------------------------------------------------------------------

[docs]
def avg_pool1d(
    input: Tensor,
    kernel_size,
    stride=None,
    padding=0,
    ceil_mode: bool = False,
    count_include_pad: bool = True,
    divisor_override=None,
) -> Tensor:
    r"""avg_pool1d(input, kernel_size, stride=None, padding=0, ceil_mode=False,
    count_include_pad=True, divisor_override=None) -> Tensor

    Applies a 1D average pooling over an input signal composed of several
    input planes. Input shape ``(N, C, L)`` or unbatched ``(C, L)``.
    """
    captured = _capture_call(avg_pool1d, (input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override), {})
    if captured is not None:
        return captured
    unbatched = input.dim() == 2
    x = input.unsqueeze(0) if unbatched else input
    k = _single(kernel_size)[0]
    s = k if stride is None else _single(stride)[0]
    p = _single(padding)[0]
    out = avg_pool2d(
        x.unsqueeze(3), (k, 1), (s, 1), (p, 0), ceil_mode,
        count_include_pad, divisor_override,
    ).squeeze(3)
    return out.squeeze(0) if unbatched else out

[docs]
def max_pool1d(
    input: Tensor,
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode: bool = False,
    return_indices: bool = False,
):
    r"""max_pool1d(input, kernel_size, stride=None, padding=0, dilation=1,
    ceil_mode=False, return_indices=False) -> Tensor

    Applies a 1D max pooling over an input signal composed of several input
    planes. Input shape ``(N, C, L)`` or unbatched ``(C, L)``.
    """
    captured = _capture_call(max_pool1d, (input, kernel_size, stride, padding, dilation, ceil_mode, return_indices), {})
    if captured is not None:
        return captured
    if return_indices:
        return max_pool1d_with_indices(
            input, kernel_size, stride=stride, padding=padding,
            dilation=dilation, ceil_mode=ceil_mode)
    unbatched = input.dim() == 2
    x = input.unsqueeze(0) if unbatched else input
    k = _single(kernel_size)[0]
    s = k if stride is None else _single(stride)[0]
    p = _single(padding)[0]
    d = _single(dilation)[0]
    out = max_pool2d(
        x.unsqueeze(3), (k, 1), (s, 1), (p, 0), (d, 1), ceil_mode,
    ).squeeze(3)
    return out.squeeze(0) if unbatched else out

[docs]
def adaptive_avg_pool1d(input: Tensor, output_size) -> Tensor:
    r"""adaptive_avg_pool1d(input, output_size) -> Tensor"""
    unbatched = input.dim() == 2
    x = input.unsqueeze(0) if unbatched else input
    out = adaptive_avg_pool2d(x.unsqueeze(3), (output_size, 1)).squeeze(3)
    return out.squeeze(0) if unbatched else out

[docs]
def adaptive_max_pool1d(input: Tensor, output_size):
    r"""adaptive_max_pool1d(input, output_size) -> Tensor"""
    unbatched = input.dim() == 2
    x = input.unsqueeze(0) if unbatched else input
    out = adaptive_max_pool2d(x.unsqueeze(3), (output_size, 1)).squeeze(3)
    return out.squeeze(0) if unbatched else out

[docs]
def lp_pool2d(
    input: Tensor,
    norm_type,
    kernel_size,
    stride=None,
    ceil_mode: bool = False,
) -> Tensor:
    r"""Apply a 2D power-average pooling over an input signal.

    If the sum of all inputs to the power of `p` is zero, the gradient is set
    to zero as well.

    See :class:`~tensorplay.nn.LPPool2d` for details.
    """
    kw, kh = _pair(kernel_size)
    if isinstance(norm_type, (int, float)):
        if norm_type == 0:
            raise ValueError(f"norm_type must be a non-zero value, but got {norm_type}")
        if norm_type == float("inf"):
            return max_pool2d(input.abs(), kernel_size, stride, 0, 1, ceil_mode)
        if norm_type == -float("inf"):
            return -max_pool2d(-input.abs(), kernel_size, stride, 0, 1, ceil_mode)

    if stride is not None:
        out = avg_pool2d(input.pow(norm_type), kernel_size, stride, 0, ceil_mode)
    else:
        out = avg_pool2d(input.pow(norm_type), kernel_size, padding=0, ceil_mode=ceil_mode)

    return (tensorplay.sign(out) * relu(tensorplay.abs(out))).mul(kw * kh).pow(1.0 / norm_type)

[docs]
def lp_pool1d(
    input: Tensor,
    norm_type,
    kernel_size,
    stride=None,
    ceil_mode: bool = False,
) -> Tensor:
    r"""Apply a 1D power-average pooling over an input signal.

    See :class:`~tensorplay.nn.LPPool1d` for details.
    """
    k = _single(kernel_size)[0]
    s = None if stride is None else _single(stride)[0]

    if isinstance(norm_type, (int, float)):
        if norm_type == 0:
            raise ValueError(f"norm_type must be a non-zero value, but got {norm_type}")
        if norm_type == float("inf"):
            return max_pool1d(input.abs(), kernel_size, stride, 0, 1, ceil_mode)
        if norm_type == -float("inf"):
            return -max_pool1d(-input.abs(), kernel_size, stride, 0, 1, ceil_mode)

    if stride is not None:
        out = avg_pool1d(input.pow(norm_type), kernel_size, stride, 0, ceil_mode)
    else:
        out = avg_pool1d(input.pow(norm_type), kernel_size, padding=0, ceil_mode=ceil_mode)

    return (tensorplay.sign(out) * relu(tensorplay.abs(out))).mul(k).pow(1.0 / norm_type)

[docs]
def local_response_norm(
    input: Tensor,
    size: int,
    alpha: float = 1e-4,
    beta: float = 0.75,
    k: float = 1.0,
) -> Tensor:
    r"""Apply local response normalization over an input signal.

    The input signal is composed of several input planes, where channels
    occupy the second dimension. Normalization is applied across channels.

    See :class:`~tensorplay.nn.LocalResponseNorm` for details.
    """
    captured = _capture_call(local_response_norm, (input, size, alpha, beta, k), {})
    if captured is not None:
        return captured
    dim = input.dim()
    if dim < 3:
        raise ValueError(
            f"Expected 3D or higher dimensionality input (got {dim} dimensions)"
        )
    if input.numel() == 0:
        return input

    # Windowed sum of squares along the channel axis (dim 1), equivalent to
    div = input.mul(input)
    pad_left = size // 2
    pad_right = (size - 1) // 2

    def _zero_channels(n: int, ref: Tensor) -> Tensor:
        shape = list(ref.shape)
        shape[1] = n
        return tensorplay.zeros(shape, dtype=ref.dtype, device=ref.device)

    # The leading zero channel beyond ``pad_left`` makes the prefix sums
    # 1-indexed, so the window covering output channel ``i`` -- padded
    # channels ``[i, i + size)`` -- is exactly ``cs[i + size] - cs[i]``.
    parts = [_zero_channels(pad_left + 1, div), div]
    if pad_right:
        parts.append(_zero_channels(pad_right, div))
    padded = tensorplay.cat(parts, dim=1)

    cs = padded.cumsum(1)
    c = input.size(1)
    hi = tensorplay.narrow(cs, 1, size, c)
    lo = tensorplay.narrow(cs, 1, 0, c)
    window_sum = hi - lo

    div = window_sum.mul(alpha / size).add(k).pow(beta)
    return input / div

[docs]
def dropout1d(
    input: Tensor,
    p: float = 0.5,
    training: bool = True,
    inplace: bool = False,
) -> Tensor:
    r"""Randomly zero out entire channels (a channel is a 1D feature map).

    See :class:`~tensorplay.nn.Dropout1d` for details.
    """
    if p < 0.0 or p > 1.0:
        raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")
    inp_dim = input.dim()
    if inp_dim not in (2, 3):
        raise RuntimeError(
            f"dropout1d: Expected 2D or 3D input, but received a {inp_dim}D input. "
            "Note that dropout1d exists to provide channel-wise dropout on inputs with 1 "
            "spatial dimension, a channel dimension, and an optional batch dimension "
            "(i.e. 2D or 3D inputs)."
        )
    if not training or p == 0:
        return input

    if p == 1:
        result = _C.zeros_like(input)
        return input.copy_(result) if inplace else result

    is_batched = inp_dim == 3
    x = input if is_batched else input.unsqueeze(0)

    mask_shape = list(x.shape)
    mask_shape[2:] = [1] * (x.dim() - 2)
    mask = (_C.rand(mask_shape, device=x.device) > p).to(x.dtype)
    scale = 1.0 / (1.0 - p)

    if inplace:
        result = x.mul_(mask).mul_(scale)
    else:
        result = x * mask * scale

    if not is_batched:
        result = result.squeeze(0)
    return result

[docs]
def feature_alpha_dropout(
    input: Tensor,
    p: float = 0.5,
    training: bool = False,
    inplace: bool = False,
) -> Tensor:
    r"""Randomly masks out entire channels, setting activations to the
    negative saturation value of the SELU activation function.

    See :class:`~tensorplay.nn.FeatureAlphaDropout` for details.
    """
    if p < 0.0 or p > 1.0:
        raise ValueError(f"dropout probability has to be between 0 and 1, but got {p}")

    if input.dim() < 2:
        raise RuntimeError(
            f"feature_alpha_dropout: Expected input to have at least 2 dimensions, "
            f"but got {input.dim()}"
        )

    if not training or p == 0 or input.numel() == 0:
        return input

    if p == 1:
        result = _C.zeros_like(input)
        return input.copy_(result) if inplace else result

    alpha_c = 1.7580993408473766
    a = 1.0 / math.sqrt((alpha_c * alpha_c * p + 1) * (1 - p))
    b_coeff = alpha_c * a

    mask_shape = list(input.shape)
    mask_shape[2:] = [1] * (input.dim() - 2)
    noise = (_C.rand(mask_shape, device=input.device) > p).to(input.dtype)
    b = (noise - 1).mul(b_coeff).add(b_coeff * p)
    noise = noise.mul(a)

    result = input * noise + b
    if inplace:
        return input.copy_(result)
    return result

# -----------------------------------------------------------------------------
# Distance helpers
# -----------------------------------------------------------------------------

_COSINE_FLOATING_DTYPES = frozenset(
    (
        tensorplay.float16,
        tensorplay.bfloat16,
        tensorplay.float32,
        tensorplay.float64,
    )
)

[docs]
def cosine_similarity(x1: Tensor, x2: Tensor, dim: int = 1, eps: float = 1e-8) -> Tensor:
    r"""Returns cosine similarity between x1 and x2, computed along dim.

    The inputs are normalized before the dot product, and the per-row norms
    are accumulated in a wider dtype (float32 for half-precision inputs,
    float64 for float32), so that squaring large values cannot overflow the
    accumulator and turn the result into NaN.
    """
    common = tensorplay.promote_types(x1.dtype, x2.dtype)
    if common not in _COSINE_FLOATING_DTYPES:
        raise RuntimeError(
            f"expected common dtype to be floating point, yet common dtype is {common}"
        )
    if x1.dtype != common:
        x1 = x1.to(common)
    if x2.dtype != common:
        x2 = x2.to(common)
    x1, x2 = tensorplay.broadcast_tensors(x1, x2)
    acc = (
        tensorplay.float32
        if common in (tensorplay.float16, tensorplay.bfloat16)
        else tensorplay.float64
    )
    n1 = x1.to(acc).norm(dim=dim, keepdim=True).clamp_min(eps).to(common)
    n2 = x2.to(acc).norm(dim=dim, keepdim=True).clamp_min(eps).to(common)
    return ((x1 / n1) * (x2 / n2)).sum(dim)

# -----------------------------------------------------------------------------
# Mask helpers used by Transformer modules.
# -----------------------------------------------------------------------------

def _canonical_mask(
    mask: Optional[Tensor],
    mask_name: str,
    other_type=None,
    other_name: str = "",
    target_type=None,
    check_other: bool = True,
) -> Optional[Tensor]:
    if mask is not None:
        _mask_dtype = mask.dtype
        _mask_is_float = _mask_dtype in (
            DType.float16, DType.bfloat16, DType.float32, DType.float64,
        )
        if _mask_dtype != DType.bool and not _mask_is_float:
            raise AssertionError(
                f"only bool and floating types of {mask_name} are supported"
            )
        if check_other and other_type is not None:
            if _mask_dtype != other_type:
                warnings.warn(
                    f"Support for mismatched {mask_name} and {other_name} "
                    "is deprecated. Use same type for both instead.",
                    stacklevel=2,
                )
        if not _mask_is_float:
            mask = tensorplay.zeros_like(mask, dtype=target_type).masked_fill_(
                mask, float("-inf")
            )
    return mask

def _none_or_dtype(input: Optional[Tensor]):
    if input is None:
        return None
    elif isinstance(input, Tensor):
        return input.dtype
    raise RuntimeError("input to _none_or_dtype() must be None or Tensor")

# -----------------------------------------------------------------------------
# live in the native schema and backend kernels).
# -----------------------------------------------------------------------------

[docs]
def multilabel_soft_margin_loss(
    input: Tensor,
    target: Tensor,
    weight: Optional[Tensor] = None,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the multilabel soft margin loss.

    See :class:`~tensorplay.nn.MultiLabelSoftMarginLoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)

    loss = -(target * logsigmoid(input) + (1 - target) * logsigmoid(-input))

    if weight is not None:
        loss = loss * weight

    class_dim = input.dim() - 1
    C = input.size(class_dim)
    loss = loss.sum(class_dim) / C  # only return N loss values

    if reduction == "none":
        ret = loss
    elif reduction == "mean":
        ret = loss.mean()
    elif reduction == "sum":
        ret = loss.sum()
    else:
        ret = input
        raise ValueError(reduction + " is not valid")
    return ret

[docs]
def gaussian_nll_loss(
    input: Tensor,
    target: Tensor,
    var,
    full: bool = False,
    eps: float = 1e-6,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the Gaussian negative log likelihood loss.

    See :class:`~tensorplay.nn.GaussianNLLLoss` for details.
    """
    # Entries of var must be non-negative
    if isinstance(var, float):
        if var < 0:
            raise ValueError("var has negative entry/entries")
        var = tensorplay.ones_like(input) * var
    elif (var < 0).to(input.dtype).sum().item() > 0:
        raise ValueError("var has negative entry/entries")

    # Check var size
    if tuple(var.size()) != tuple(input.size()):
        # If var is one dimension short of input, but the sizes match otherwise,
        # then this is a homoscedastic case.
        if tuple(input.size())[:-1] == tuple(var.size()):
            var = tensorplay.unsqueeze(var, -1)
        elif (
            input.dim() == var.dim()
            and sum(y for x, y in zip(input.size(), var.size()) if x != y) == 1
        ):  # Heteroscedastic case
            pass
        else:
            raise ValueError("var is of incorrect size")

    # Check validity of reduction mode
    if reduction not in ("mean", "sum", "none"):
        raise ValueError(reduction + " is not valid")

    # Clamp for stability
    var = var.clone()
    with tensorplay.no_grad():
        var.clamp_(min=eps)

    # Calculate the loss
    loss = 0.5 * (tensorplay.log(var) + (input - target) ** 2 / var)
    if full:
        loss += 0.5 * math.log(2 * math.pi)

    if reduction == "mean":
        return loss.mean()
    elif reduction == "sum":
        return loss.sum()
    else:
        return loss

[docs]
def triplet_margin_with_distance_loss(
    anchor: Tensor,
    positive: Tensor,
    negative: Tensor,
    *,
    distance_function=None,
    margin: float = 1.0,
    swap: bool = False,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the triplet margin loss using a custom distance function.

    See :class:`~tensorplay.nn.TripletMarginWithDistanceLoss` for details.
    """
    # Check validity of reduction mode
    if reduction not in ("mean", "sum", "none"):
        raise ValueError(f"{reduction} is not a valid value for reduction")

    # Check validity of margin
    if margin <= 0:
        raise ValueError(f"margin must be greater than 0, got {margin}")

    # Check dimensions
    a_dim = anchor.dim()
    p_dim = positive.dim()
    n_dim = negative.dim()
    if not (a_dim == p_dim and p_dim == n_dim):
        raise RuntimeError(
            f"The anchor, positive, and negative tensors are expected to have "
            f"the same number of dimensions, but got: anchor {a_dim}D, "
            f"positive {p_dim}D, and negative {n_dim}D inputs"
        )

    # Calculate loss
    if distance_function is None:
        distance_function = tensorplay.pairwise_distance

    dist_pos = distance_function(anchor, positive)
    dist_neg = distance_function(anchor, negative)
    # The distance swap is described in the paper "Learning shallow
    # convolutional feature descriptors with triplet losses" by V. Balntas, E.
    # Riba et al.
    if swap:
        dist_swap = distance_function(positive, negative)
        dist_neg = tensorplay.minimum(dist_neg, dist_swap)
    loss = tensorplay.clamp(margin + dist_pos - dist_neg, 0.0)

    # Apply reduction
    if reduction == "sum":
        return loss.sum()
    elif reduction == "mean":
        return loss.mean()
    else:  # reduction == "none"
        return loss

# =============================================================================
# 2.15.0a0 @ 893b6406).  Ops that are native in this repo call the dispatcher
# directly; the rest are composed from dispatched primitives following the
# same math without new kernels.
# =============================================================================

def _broadcast_shapes(*shapes):
    out = []
    for dims in zip(*(reversed(s) for s in shapes)):
        d = 1
        for x in dims:
            if x != 1:
                if d != 1 and x != d:
                    raise ValueError("Shape mismatch: objects cannot be broadcast to a single shape")
                d = x
        out.append(d)
    return tuple(reversed(out))

def _apply_reduction(loss, reduction):
    if reduction == "none":
        return loss
    if reduction == "mean":
        return loss.mean()
    if reduction == "sum":
        return loss.sum()
    raise ValueError(f"{reduction} is not a valid value for reduction")

def _band(t, lo, hi):
    """t >= lo AND t < hi, elementwise (bool)."""
    return tensorplay.logical_and(t >= lo, t < hi)

[docs]
def one_hot(tensor: Tensor, num_classes: int = -1) -> Tensor:
    r"""one_hot(tensor, num_classes=-1) -> LongTensor

    Returns long tensor shaped ``tensor.shape + (num_classes,)`` with a 1 at
    """
    captured = _capture_call(one_hot, (tensor, num_classes), {})
    if captured is not None:
        return captured
    if num_classes < 0:
        if tensor.numel() == 0:
            raise RuntimeError("Cannot infer num classes from empty tensor")
        num_classes = int(tensor.max().item()) + 1
    t64 = tensor.to(DType.int64)
    rng = tensorplay.arange(num_classes, dtype=DType.int64, device=tensor.device)
    return t64.unsqueeze(-1).eq(rng).to(DType.int64)

[docs]
def sigmoid(input):
    r"""sigmoid(input) -> Tensor

    Applies the element-wise function :math:`\text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)}`
    """
    return tensorplay.sigmoid(input)

[docs]
def tanh(input):
    r"""tanh(input) -> Tensor

    Applies element-wise :math:`\text{Tanh}(x) = \frac{\exp(x) - \exp(-x)}{\exp(x) + \exp(-x)}`
    """
    return tensorplay.tanh(input)

[docs]
def rms_norm(
    input: Tensor,
    normalized_shape,
    weight: Optional[Tensor] = None,
    eps: Optional[float] = None,
) -> Tensor:
    r"""Apply Root Mean Square Layer Normalization.

    Dispatches to the native fused kernel (single dispatch, CPU vectorized
    rows / CUDA block-per-row); falls back to the composite below under
    CompositeImplicitAutograd rms_norm."""
    return _rms_norm_impl(input, normalized_shape, weight, eps)

def _rms_norm_composite(
    input: Tensor,
    normalized_shape,
    weight: Optional[Tensor] = None,
    eps: Optional[float] = None,
) -> Tensor:
    r"""
    rms_norm composite (fp32 compute for reduced dtypes)."""
    shape = list(normalized_shape) if isinstance(normalized_shape, (list, tuple)) else [int(normalized_shape)]
    ndim = len(shape)
    dims = list(range(input.dim() - ndim, input.dim()))
    if eps is None:
        eps = 1e-5
    compute_dtype = DType.float32 if input.dtype in (DType.float16, DType.bfloat16) else input.dtype
    x = input.to(compute_dtype)
    denom = x.pow(2).mean(dims, keepdim=True) + eps
    inv = tensorplay.rsqrt(denom)
    out = x * inv
    if weight is not None:
        w_shape = [1] * (input.dim() - ndim) + list(shape)
        out = out * weight.to(compute_dtype).view(w_shape)
    return out.to(input.dtype)

def _rms_norm_impl(
    input: Tensor,
    normalized_shape,
    weight: Optional[Tensor] = None,
    eps: Optional[float] = None,
) -> Tensor:
    """Native fused kernel (CPU vectorized rows / CUDA block-per-row).

    Falls back to the composite above under autograd: the native forward has
    flows through its inner ops), so training graphs must keep composing."""
    needs_grad = tensorplay.is_grad_enabled() and (
        input.requires_grad
        or (weight is not None and getattr(weight, "requires_grad", False))
    )
    if needs_grad:
        return _rms_norm_composite(input, normalized_shape, weight, eps)
    return tensorplay._C.rms_norm(
        input,
        list(normalized_shape) if isinstance(normalized_shape, (list, tuple)) else [int(normalized_shape)],
        weight,
        float(eps) if eps is not None else None,
    )

[docs]
def gumbel_softmax(
    logits: Tensor,
    tau: float = 1,
    hard: bool = False,
    eps: float = 1e-10,
    dim: int = -1,
) -> Tensor:
    r"""Sample from the Gumbel-Softmax distribution and optionally discretize.

    straight-through when ``hard=True``).
    """
    if eps != 1e-10:
        warnings.warn("`eps` parameter is deprecated and has no effect.", stacklevel=2)
    u = tensorplay.empty_like(logits)
    u.exponential_()
    gumbels = -u.log()
    gumbels = (logits + gumbels) / tau
    y_soft = softmax(gumbels, dim=dim)
    if hard:
        idx = _C.argmax(y_soft, dim=dim, keepdim=True)
        ar = tensorplay.arange(y_soft.size(dim), dtype=DType.int64, device=logits.device)
        view = [1] * logits.dim()
        view[dim] = -1
        y_hard = idx.eq(ar.view(view)).to(logits.dtype)
        return y_hard.detach() - y_soft.detach() + y_soft
    return y_soft

# -----------------------------------------------------------------------------
# Loss functions (F-alignment)
# -----------------------------------------------------------------------------

[docs]
def l1_loss(
    input: Tensor,
    target: Tensor,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
    weight: Optional[Tensor] = None,
) -> Tensor:
    r"""Compute the L1 loss, with optional weighting.

    Function that takes the mean element-wise absolute value difference.
    See :class:`~tensorplay.nn.L1Loss` for details.
    """
    captured = _capture_call(l1_loss, (input, target, size_average, reduce, reduction, weight), {})
    if captured is not None:
        return captured
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)
    if target.size() != input.size():
        warnings.warn(
            f"Using a target size ({target.size()}) that is different to the input size ({input.size()}). "
            "This will likely lead to incorrect results due to broadcasting. "
            "Please ensure they have the same size.",
            stacklevel=2,
        )
    if weight is not None:
        if weight.size() != input.size():
            raise ValueError("Weights and input must have the same size.")
        absolute_errors = (input - target).abs()
        weighted = absolute_errors * weight
        if reduction == "none":
            return weighted
        if reduction == "sum":
            return weighted.sum()
        return weighted.sum() / weight.sum()

    expanded_input, expanded_target = _expand_pair(input, target)
    return _C.tp_l1_loss(expanded_input, expanded_target, _get_reduction_enum(reduction))

def _expand_pair(input: Tensor, target: Tensor):
    """Broadcast input/target to a common shape when they differ."""
    if tuple(input.size()) == tuple(target.size()):
        return input, target
    shape = _broadcast_shapes(tuple(input.size()), tuple(target.size()))
    return input.expand(shape), target.expand(shape)

_REDUCTION_STRINGS = {0: "none", 1: "mean", 2: "sum"}

[docs]
def smooth_l1_loss(
    input: Tensor,
    target: Tensor,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
    beta: float = 1.0,
) -> Tensor:
    r"""Compute the Smooth L1 loss.

    Function uses a squared term if the absolute element-wise error falls
    below beta and an L1 term otherwise.
    See :class:`~tensorplay.nn.SmoothL1Loss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)
    if target.size() != input.size():
        warnings.warn(
            f"Using a target size ({target.size()}) that is different to the input size ({input.size()}). "
            "This will likely lead to incorrect results due to broadcasting. "
            "Please ensure they have the same size.",
            stacklevel=2,
        )
    expanded_input, expanded_target = _expand_pair(input, target)
    enum_ = _get_reduction_enum(reduction)
    if beta == 0.0:
        return _C.tp_l1_loss(expanded_input, expanded_target, enum_)
    return _C.smooth_l1_loss(expanded_input, expanded_target, enum_, beta)

[docs]
def huber_loss(
    input: Tensor,
    target: Tensor,
    reduction: str = "mean",
    delta: float = 1.0,
    weight: Optional[Tensor] = None,
) -> Tensor:
    r"""Compute the Huber loss, with optional weighting.

    Function uses a squared term if the absolute error falls below delta and
    a delta-scaled L1 term otherwise.
    See :class:`~tensorplay.nn.HuberLoss` for details.
    """
    if target.size() != input.size():
        warnings.warn(
            f"Using a target size ({target.size()}) that is different to the input size ({input.size()}). "
            "This will likely lead to incorrect results due to broadcasting. "
            "Please ensure they have the same size.",
            stacklevel=2,
        )
    expanded_input, expanded_target = _expand_pair(input, target)
    enum_ = _get_reduction_enum(reduction)
    if weight is None:
        return _C.huber_loss(expanded_input, expanded_target, enum_, delta)

    if weight.size() != input.size():
        raise ValueError("Weights and input must have the same size.")
    unweighted = _C.huber_loss(expanded_input, expanded_target, 0, delta)
    weighted = unweighted * weight
    if reduction == "none":
        return weighted
    if reduction == "sum":
        return weighted.sum()
    return weighted.mean()

[docs]
def kl_div(
    input: Tensor,
    target: Tensor,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
    log_target: bool = False,
) -> Tensor:
    r"""Compute the KL Divergence loss.

    ``input`` holds log-probabilities; see :class:`~tensorplay.nn.KLDivLoss`.
    Note that :attr:`reduction='mean'` divides by the number of elements and
    does not return the true KL divergence value — use ``'batchmean'``.
    """
    if size_average is not None or reduce is not None:
        reduction_enum = _get_reduction_enum(_legacy_get_string(size_average, reduce))
    else:
        if reduction == "mean":
            warnings.warn(
                "reduction: 'mean' divides the total loss by both the batch size and the support size."
                "'batchmean' divides only by the batch size, and aligns with the KL div math definition."
                "'mean' will be changed to behave the same as 'batchmean' in the next major release.",
                stacklevel=2,
            )
        if reduction == "batchmean":
            reduction_enum = 2  # sum
        else:
            reduction_enum = _get_reduction_enum(reduction)

    reduced = _C.tp_kl_div(input, target, reduction_enum, log_target)
    if reduction == "batchmean" and input.dim() != 0:
        reduced = reduced / input.size(0)
    return reduced

[docs]
def binary_cross_entropy(
    input: Tensor,
    target: Tensor,
    weight: Optional[Tensor] = None,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute Binary Cross Entropy between the target and input probabilities.

    See :class:`~tensorplay.nn.BCELoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction_enum = _get_reduction_enum(_legacy_get_string(size_average, reduce))
    else:
        reduction_enum = _get_reduction_enum(reduction)
    if target.size() != input.size():
        raise ValueError(
            f"Using a target size ({target.size()}) that is different to the input size ({input.size()}) is deprecated. "
            "Please ensure they have the same size."
        )
    if weight is not None:
        new_size = _broadcast_shapes(tuple(target.size()), tuple(weight.size()))
        weight = weight.expand(new_size)
    return _C.binary_cross_entropy(input, target, weight, reduction_enum)

[docs]
def binary_cross_entropy_with_logits(
    input: Tensor,
    target: Tensor,
    weight: Optional[Tensor] = None,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
    pos_weight: Optional[Tensor] = None,
) -> Tensor:
    r"""Compute Binary Cross Entropy between target and input logits.

    optionally rescaled by ``weight``, then reduced.
    See :class:`~tensorplay.nn.BCEWithLogitsLoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction_enum = _get_reduction_enum(_legacy_get_string(size_average, reduce))
    else:
        reduction_enum = _get_reduction_enum(reduction)
    if tuple(target.size()) != tuple(input.size()):
        raise ValueError(
            f"Target size ({target.size()}) must be the same as input size ({input.size()})"
        )

    log_sigmoid_input = logsigmoid(input)
    if pos_weight is not None:
        log_weight = (pos_weight - 1).mul(target).add(1)
        log_sigmoid_input = log_sigmoid_input.mul(log_weight)

    loss = (1 - target).mul(input).sub(log_sigmoid_input)
    if weight is not None:
        loss = loss.mul(weight)
    return _apply_reduction(loss, _REDUCTION_STRINGS[reduction_enum])

[docs]
def poisson_nll_loss(
    input: Tensor,
    target: Tensor,
    log_input: bool = True,
    full: bool = False,
    size_average=None,
    eps: float = 1e-8,
    reduce=None,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the Poisson negative log likelihood loss.

    See :class:`~tensorplay.nn.PoissonNLLLoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)
    if reduction not in ("none", "mean", "sum"):
        raise ValueError(f"{reduction} is not a valid value for reduction")
    return _C.tp_poisson_nll_loss(
        input, target, log_input, full, eps, _get_reduction_enum(reduction)
    )

[docs]
def soft_margin_loss(
    input: Tensor,
    target: Tensor,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the soft margin loss.

    See :class:`~tensorplay.nn.SoftMarginLoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)
    return _C.tp_soft_margin_loss(input, target, _get_reduction_enum(reduction))

[docs]
def cosine_embedding_loss(
    input1: Tensor,
    input2: Tensor,
    target: Tensor,
    margin: float = 0,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the cosine embedding loss.

    See :class:`~tensorplay.nn.CosineEmbeddingLoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)
    return _C.tp_cosine_embedding_loss(
        input1, input2, target, margin, _get_reduction_enum(reduction)
    )

[docs]
def margin_ranking_loss(
    input1: Tensor,
    input2: Tensor,
    target: Tensor,
    margin: float = 0,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the margin ranking loss.

    See :class:`~tensorplay.nn.MarginRankingLoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)
    return _C.tp_margin_ranking_loss(
        input1, input2, target, margin, _get_reduction_enum(reduction)
    )

[docs]
def hinge_embedding_loss(
    input: Tensor,
    target: Tensor,
    margin: float = 1.0,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the hinge embedding loss.

    See :class:`~tensorplay.nn.HingeEmbeddingLoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)
    return _C.tp_hinge_embedding_loss(
        input, target, margin, _get_reduction_enum(reduction)
    )

[docs]
def multi_margin_loss(
    input: Tensor,
    target: Tensor,
    p: int = 1,
    margin: float = 1.0,
    weight: Optional[Tensor] = None,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the multi margin loss, with optional weighting.

    ``sum_d max(0, margin - x_y + x_d)^p * w_y / C`` over non-target classes.
    See :class:`~tensorplay.nn.MultiMarginLoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)
    if p != 1 and p != 2:
        raise ValueError("only p == 1 and p == 2 supported")
    if weight is not None and weight.dim() != 1:
        raise ValueError("weight must be one-dimensional")
    return _C.multi_margin_loss(input, target.to(DType.int64), p, margin, weight,
                                _get_reduction_enum(reduction))

[docs]
def multilabel_margin_loss(
    input: Tensor,
    target: Tensor,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the multilabel margin loss.

    for each positive label ``y`` (targets are active until the first
    ``-1``), add ``max(0, 1 - x[y] + x[d])`` over non-target labels ``d``;
    divide by C.
    See :class:`~tensorplay.nn.MultiLabelMarginLoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)
    return _C.multilabel_margin_loss(input, target.to(DType.int64),
                                     _get_reduction_enum(reduction))

[docs]
def triplet_margin_loss(
    anchor: Tensor,
    positive: Tensor,
    negative: Tensor,
    margin: float = 1.0,
    p: float = 2,
    eps: float = 1e-6,
    swap: bool = False,
    size_average=None,
    reduce=None,
    reduction: str = "mean",
) -> Tensor:
    r"""Compute the triplet loss between given input tensors and a margin
    greater than 0.

    See :class:`~tensorplay.nn.TripletMarginLoss` for details.
    """
    if size_average is not None or reduce is not None:
        reduction = _legacy_get_string(size_average, reduce)
    if margin <= 0:
        raise ValueError(f"margin must be greater than 0, got {margin}")
    dist_pos = pairwise_distance(anchor, positive, p, eps)
    dist_neg = pairwise_distance(anchor, negative, p, eps)
    # The distance swap is described in the paper "Learning shallow
    # convolutional feature descriptors with triplet losses" (Balntas et al.).
    if swap:
        dist_swap = pairwise_distance(positive, negative, p, eps)
        dist_neg = tensorplay.minimum(dist_neg, dist_swap)
    loss = tensorplay.clamp(margin + dist_pos - dist_neg, min=0.0)
    return _apply_reduction(loss, reduction)

[docs]
def ctc_loss(
    log_probs: Tensor,
    targets: Tensor,
    input_lengths,
    target_lengths,
    blank: int = 0,
    reduction: str = "mean",
    zero_infinity: bool = False,
) -> Tensor:
    r"""Compute the Connectionist Temporal Classification loss.

    (alpha recurrence over the blank-extended target sequence); autograd flows
    through ``log_probs`` via the composed primitives.

    Args:
        log_probs: :math:`(T, N, C)` or :math:`(T, C)` log-softmax outputs.
        targets: :math:`(N, S)` or concatenated :math:`(\sum S_n,)`.
        input_lengths / target_lengths: :math:`(N,)` or scalars.
        blank: index of the blank label. Default: 0.
        reduction: ``'none' | 'mean' | 'sum'``.
        zero_infinity: zero out infinite losses (targets too long for T).
    """
    captured = _capture_call(
        ctc_loss,
        (log_probs, targets, input_lengths, target_lengths, blank, reduction,
         zero_infinity),
        {},
    )
    if captured is not None:
        return captured
    if log_probs.dim() == 2:
        log_probs = log_probs.unsqueeze(1)
        unbatched = True
    else:
        unbatched = False
    if log_probs.dim() != 3:
        raise ValueError(f"ctc_loss: log_probs must be 2D or 3D, got {log_probs.dim()}D")
    T, N, C = log_probs.shape

    def _as_len_tensor(x):
        if x is None:
            return None
        if isinstance(x, Tensor):
            return x.to(DType.int64).reshape(-1)
        if isinstance(x, int):
            return tensorplay.full([N], x, dtype=DType.int64, device=log_probs.device)
        return tensorplay.tensor(list(x), dtype=DType.int64, device=log_probs.device)

    in_lens = _as_len_tensor(input_lengths)
    tgt_lens = _as_len_tensor(target_lengths)

    if targets.dim() == 1:
        rows = []
        off = 0
        for n in range(N):
            s = int(tgt_lens[n].item())
            rows.append(tensorplay.narrow(targets.to(DType.int64), 0, off, s))
            off += s
        S = max((int(r.numel()) for r in rows), default=0)
        padded = tensorplay.zeros(N, S, dtype=DType.int64, device=log_probs.device)
        for n in range(N):
            k = rows[n].numel()
            if k:
                padded[n, :k] = rows[n]
        targets2d = padded
    elif targets.dim() == 2:
        targets2d = targets.to(DType.int64)
        S = targets2d.size(1)
    else:
        raise ValueError(f"ctc_loss: targets must be 1D or 2D, got {targets.dim()}D")

    # Native dispatch: _ctc_loss + reduction compose on the dispatcher, and
    # the _ctc_loss derivative formula (derivatives.yaml) drives the backward.
    return _ctc_loss_impl(log_probs, targets2d, in_lens, tgt_lens, blank,
                          reduction, zero_infinity, unbatched)

def _ctc_zero_inf_mask(nll):
    # impossible alignments carry +inf (or NaN from inf - inf) raw NLL
    return tensorplay.logical_or(nll.ge(float("inf")), nll.ne(nll))

def _ctc_loss_impl(log_probs, targets2d, in_lens, tgt_lens, blank, reduction,
                   zero_infinity, unbatched):
    # NB: the FASTCALL binding layer for underscore ops requires kwargs.
    nll, _ = _C._ctc_loss(log_probs=log_probs, targets=targets2d,
                          input_lengths=in_lens, target_lengths=tgt_lens,
                          blank=blank, zero_infinity=zero_infinity)
    if zero_infinity:
        nll = tensorplay.where(_ctc_zero_inf_mask(nll),
                               tensorplay.zeros_like(nll), nll)
    if reduction == "none":
        return nll.squeeze(0) if unbatched else nll
    if reduction == "sum":
        return nll.sum()
    # 'mean': divide each by its target length, then average
    denom = tgt_lens.to(nll.dtype).clamp(min=1)
    return (nll / denom).mean()

[docs]
def pixel_shuffle(input: Tensor, upscale_factor: int) -> Tensor:
    r"""Rearranges elements in a tensor of shape ``(*, C x r^2, H, W)`` to a
    tensor of shape ``(*, C, H x r, W x r)``.

    input[n, c*r^2 + i*r + j, h, w]``.
    """
    captured = _capture_call(pixel_shuffle, (input, upscale_factor), {})
    if captured is not None:
        return captured
    r = int(upscale_factor)
    if input.dim() != 4:
        raise ValueError(f"pixel_shuffle expects 4D input, got {input.dim()}D")
    N, C, H, W = input.shape
    if C % (r * r) != 0:
        raise ValueError(
            f"pixel_shuffle expects input channel to be divisible by square of upscale_factor, "
            f"but got C={C} and upscale_factor={r}")
    Cc = C // (r * r)
    x = input.reshape(N, Cc, r, r, H, W)
    x = x.permute(0, 1, 4, 2, 5, 3).contiguous()
    return x.reshape(N, Cc, H * r, W * r)

[docs]
def pixel_unshuffle(input: Tensor, downscale_factor: int) -> Tensor:
    r"""Reverses the :func:`pixel_shuffle` transformation: ``(*, C, H x r,
    W x r) -> (*, C x r^2, H, W)``."""
    captured = _capture_call(pixel_unshuffle, (input, downscale_factor), {})
    if captured is not None:
        return captured
    r = int(downscale_factor)
    if input.dim() != 4:
        raise ValueError(f"pixel_unshuffle expects 4D input, got {input.dim()}D")
    N, C, H, W = input.shape
    if H % r != 0 or W % r != 0:
        raise ValueError(
            f"pixel_unshuffle expects input height and width divisible by downscale_factor, "
            f"but got H={H}, W={W} and downscale_factor={r}")
    x = input.reshape(N, C, H // r, r, W // r, r)
    x = x.permute(0, 1, 3, 5, 2, 4).contiguous()
    return x.reshape(N, C * r * r, H // r, W // r)

[docs]
def channel_shuffle(input: Tensor, groups: int) -> Tensor:
    r"""Divide the channels in a tensor into ``g`` groups and rearrange them
    as in ShuffleNet: ``(*, C, H, W) -> (*, C, H, W)`` with channels
    interleaved across groups."""
    g = int(groups)
    if input.dim() != 4:
        raise ValueError(f"channel_shuffle expects 4D input, got {input.dim()}D")
    N, C, H, W = input.shape
    if C % g != 0:
        raise ValueError(f"channel_shuffle expects channel count divisible by groups, got C={C}, groups={g}")
    x = input.reshape(N, g, C // g, H, W)
    x = x.transpose(1, 2).contiguous()
    return x.reshape(N, C, H, W)

[docs]
def native_channel_shuffle(input: Tensor, groups: int) -> Tensor:
    return channel_shuffle(input, groups)

GRID_SAMPLE_INTERPOLATION_MODES = ('bilinear', 'nearest', 'bicubic')
GRID_SAMPLE_PADDING_MODES = ('zeros', 'border', 'reflection')

def _linspace_from_neg_one(num_steps, align_corners, dtype, device):
    """Build normalized sampling coordinates for an affine grid."""
    if num_steps <= 1:
        return tensorplay.zeros([num_steps], dtype=dtype, device=device)
    rng = tensorplay.linspace(-1, 1, num_steps, dtype=dtype, device=device)
    if not align_corners:
        rng = rng * (num_steps - 1) / num_steps
    return rng

[docs]
def affine_grid(theta: Tensor, size, align_corners=None) -> Tensor:
    r"""Generate 2D or 3D flow field (sampling grid), given a batch of affine
    matrices :attr:`theta`.

    """
    if align_corners is None:
        warnings.warn(
            "Default grid_sample and affine_grid behavior has changed "
            "to align_corners=False since 1.3.0. Please specify "
            "align_corners=True if the old behavior is desired. "
            "See the documentation of affine_grid for details.",
            stacklevel=2)
        align_corners = False
    if len(size) == 4:
        n, _, h, w = size
        if theta.dim() != 3 or theta.size(0) != n or theta.size(1) != 2 or theta.size(2) != 3:
            raise ValueError(f"affine_grid: expected theta of shape ({n}, 2, 3), got {tuple(theta.shape)}")
        dtype, dev = theta.dtype, theta.device
        xs = _linspace_from_neg_one(w, align_corners, dtype, dev).view([1, 1, w]).expand(n, h, w)
        ys = _linspace_from_neg_one(h, align_corners, dtype, dev).view([1, h, 1]).expand(n, h, w)
        ones = tensorplay.ones([n, h, w], dtype=dtype, device=dev)
        flat = tensorplay.stack([xs, ys, ones], dim=-1).reshape(n, h * w, 3)
        grid = flat.bmm(theta.transpose(1, 2))
        return grid.reshape(n, h, w, 2)
    elif len(size) == 5:
        n, _, d, h, w = size
        if theta.dim() != 3 or theta.size(0) != n or theta.size(1) != 3 or theta.size(2) != 4:
            raise ValueError(f"affine_grid: expected theta of shape ({n}, 3, 4), got {tuple(theta.shape)}")
        dtype, dev = theta.dtype, theta.device
        xs = _linspace_from_neg_one(w, align_corners, dtype, dev).view([1, 1, 1, w]).expand(n, d, h, w)
        ys = _linspace_from_neg_one(h, align_corners, dtype, dev).view([1, 1, h, 1]).expand(n, d, h, w)
        zs = _linspace_from_neg_one(d, align_corners, dtype, dev).view([1, d, 1, 1]).expand(n, d, h, w)
        ones = tensorplay.ones([n, d, h, w], dtype=dtype, device=dev)
        flat = tensorplay.stack([xs, ys, zs, ones], dim=-1).reshape(n, d * h * w, 4)
        grid = flat.bmm(theta.transpose(1, 2))
        return grid.reshape(n, d, h, w, 3)
    raise ValueError("affine_grid: size must be length 4 or 5")

def _gs_unnormalize(coord, size, align_corners):
    if align_corners:
        return ((coord + 1) / 2) * (size - 1)
    return ((coord + 1) * size - 1) / 2

def _gs_adjust(coord, size, padding_mode, align_corners=False):
    if isinstance(padding_mode, str):
        padding_mode = GRID_SAMPLE_PADDING_MODES.index(padding_mode)
    if padding_mode == 1:  # border
        return coord.clamp(0, size - 1)
    if padding_mode == 2:  # reflection
        # reflect_coordinates over [twice_low/2, twice_high/2], then border clip.
        if align_corners:
            twice_low, twice_high = 0, 2 * (size - 1)
        else:
            twice_low, twice_high = -1, 2 * size - 1
        min_ = twice_low / 2
        span = (twice_high - twice_low) / 2
        if twice_low == twice_high:
            return tensorplay.zeros_like(coord)
        c = coord - min_
        c = tensorplay.where(c < 0, -c, c)
        flips = tensorplay.floor(c / span)
        extra = c - flips * span  # == fmod(c, span); sign-safe since c >= 0
        even = (tensorplay.floor(flips / 2) * 2) == flips
        out = tensorplay.where(even, extra + min_, span - extra + min_)
        return out.clamp(0, size - 1)
    return coord

def _grid_sample_gather(input, xs, ys, in_bounds):
    """Gather input planes at integer pixel coords.

    input: (N, C, H, W); xs/ys: (N, Ho, Wo) int64; in_bounds: bool or None.
    Returns (N, C, Ho, Wo) values zeroed outside bounds when requested.
    """
    N, C, H, W = input.shape
    dev = input.device
    i64 = DType.int64
    xs_c = xs.clamp(0, W - 1)
    ys_c = ys.clamp(0, H - 1)
    pos = ys_c * W + xs_c                                  # (N, Ho, Wo)
    if in_bounds is not None:
        pos = tensorplay.where(in_bounds, pos, pos * 0)
    # keep every operand rank-4: our broadcast kernel requires equal ranks.
    bidx = tensorplay.arange(N, dtype=i64, device=dev).view([N, 1, 1, 1]) * (C * H * W)
    cidx = tensorplay.arange(C, dtype=i64, device=dev).view([1, C, 1, 1]) * (H * W)
    gid = bidx + cidx + pos.unsqueeze(1)                   # (N, C, Ho, Wo)
    # NB: use reshape, not view - our current view() binding does not record
    # autograd, silently detaching the graph (reshape does).
    vals = tensorplay.embedding(input.contiguous().reshape(-1), gid.reshape(-1)).reshape(
        [N, C, xs.shape[-2], xs.shape[-1]])
    if in_bounds is not None:
        # expand (not broadcast): our engine lacks the sum-to-shape reduction
        # propagate wrong-shaped grads to non-leaf operands mid-graph.
        vals = vals * in_bounds.unsqueeze(1).expand(
            [N, C, xs.shape[-2], xs.shape[-1]]).to(vals.dtype)
    return vals

def _grid_sampler_2d(input, grid, interpolation_mode, padding_mode, align_corners):
    if isinstance(padding_mode, str):
        padding_mode = GRID_SAMPLE_PADDING_MODES.index(padding_mode)
    N, C, H_in, W_in = input.shape
    H_out, W_out = grid.shape[1], grid.shape[2]
    x = _gs_unnormalize(grid[..., 0], W_in, align_corners)
    y = _gs_unnormalize(grid[..., 1], H_in, align_corners)
    # padding-adjusted individually instead (GridSamplerKernel.cpp
    # Bicubic::get_value_bounded -> compute_coordinates).
    if padding_mode != 0 and interpolation_mode != 2:
        x = _gs_adjust(x, W_in, padding_mode, align_corners)
        y = _gs_adjust(y, H_in, padding_mode, align_corners)

    if interpolation_mode == 1:  # nearest
        xi = tensorplay.round(x).to(DType.int64)
        yi = tensorplay.round(y).to(DType.int64)
        ib = tensorplay.logical_and(
            tensorplay.logical_and(xi >= 0, xi < W_in),
            tensorplay.logical_and(yi >= 0, yi < H_in)) if padding_mode == 0 else None
        vals = _grid_sample_gather(input, xi, yi, ib)
        if x.requires_grad:
            # composite has no such automatic path (embedding backward emits
            # gradients for weights only), so bridge an exact-zero term.
            vals = vals + (x * 0).sum()
        return vals

    x0f = tensorplay.floor(x)
    y0f = tensorplay.floor(y)
    x0 = x0f.to(DType.int64)
    y0 = y0f.to(DType.int64)
    wx = x - x0f
    wy = y - y0f

    def corner_mask(cx, cy):
        return tensorplay.logical_and(
            tensorplay.logical_and(cx >= 0, cx < W_in),
            tensorplay.logical_and(cy >= 0, cy < H_in))

    if interpolation_mode == 0:  # bilinear
        corners = [
            (x0, y0, (1 - wx) * (1 - wy)),
            (x0 + 1, y0, wx * (1 - wy)),
            (x0, y0 + 1, (1 - wx) * wy),
            (x0 + 1, y0 + 1, wx * wy),
        ]
        out = None
        for cx, cy, wgt in corners:
            ib = corner_mask(cx, cy) if padding_mode == 0 else None
            v = _grid_sample_gather(input, cx, cy, ib)
            wgt = wgt.unsqueeze(1).expand([N, C, H_out, W_out]).to(v.dtype)
            term = v * wgt
            out = term if out is None else out + term
        return out

    # Bicubic interpolation with Keys parameter alpha = -0.75.
    def _c1(t_, a=-0.75):
        return ((a + 2) * t_ - (a + 3)) * t_ * t_ + 1

    def _c2(t_, a=-0.75):
        return ((a * t_ - 5 * a) * t_ + 8 * a) * t_ - 4 * a

    tx = x - x0f
    ty = y - y0f
    wxs = [_c2(tx + 1), _c1(tx), _c1(1 - tx), _c2(2 - tx)]
    wys = [_c2(ty + 1), _c1(ty), _c1(1 - ty), _c2(2 - ty)]
    out = None
    for j in range(4):          # rows (y offset j-1)
        for k in range(4):      # cols (x offset k-1)
            cxf = x0f + (k - 1)
            cyf = y0f + (j - 1)
            if padding_mode == 2:
                # reflection: adjust each tap index (border is handled by the
                # clamp inside _grid_sample_gather; zeros masks per tap).
                cxf = _gs_adjust(cxf, W_in, 2, align_corners)
                cyf = _gs_adjust(cyf, H_in, 2, align_corners)
            cx = cxf.to(DType.int64)
            cy = cyf.to(DType.int64)
            ib = corner_mask(cx, cy) if padding_mode == 0 else None
            v = _grid_sample_gather(input, cx, cy, ib)
            wgt = (wxs[k] * wys[j]).unsqueeze(1).expand([N, C, H_out, W_out]).to(v.dtype)
            term = v * wgt
            out = term if out is None else out + term
    return out

[docs]
def grid_sample(
    input: Tensor,
    grid: Tensor,
    mode: str = "bilinear",
    padding_mode: str = "zeros",
    align_corners=None,
) -> Tensor:
    r"""Compute grid sample.

    Given an :attr:`input` and a flow-field :attr:`grid`, computes the
    ``output`` using :attr:`input` values and pixel locations from
    :attr:`grid`. Currently, only spatial (4-D) and volumetric (5-D)
    :attr:`input` are supported.

    Args:
        input (Tensor): input of shape :math:`(N, C, H_\text{in}, W_\text{in})` (4-D case)
                        or :math:`(N, C, D_\text{in}, H_\text{in}, W_\text{in})` (5-D case)
        grid (Tensor): flow-field of shape :math:`(N, H_\text{out}, W_\text{out}, 2)` (4-D case)
                       or :math:`(N, D_\text{out}, H_\text{out}, W_\text{out}, 3)` (5-D case)
        mode (str): ``'bilinear'`` | ``'nearest'`` | ``'bicubic'``. Default: ``'bilinear'``
        padding_mode (str): ``'zeros'`` | ``'border'`` | ``'reflection'``. Default: ``'zeros'``
        align_corners (bool, optional): extrema treatment, default ``False``.

    Dispatches to the native grid_sampler_2d / grid_sampler_3d kernels
    both :attr:`input` and :attr:`grid`.
    """
    if mode not in GRID_SAMPLE_INTERPOLATION_MODES:
        raise ValueError(
            f"nn.functional.grid_sample(): expected mode to be 'bilinear', 'nearest' or 'bicubic', but got: '{mode}'")
    if padding_mode not in GRID_SAMPLE_PADDING_MODES:
        raise ValueError(
            "nn.functional.grid_sample(): expected padding_mode "
            "to be 'zeros', 'border', or 'reflection', "
            f"but got: '{padding_mode}'")
    if align_corners is None:
        warnings.warn(
            "Default grid_sample and affine_grid behavior has changed "
            "to align_corners=False since 1.3.0. Please specify "
            "align_corners=True if the old behavior is desired. "
            "See the documentation of grid_sample for details.",
            stacklevel=2)
        align_corners = False
    mode_enum = GRID_SAMPLE_INTERPOLATION_MODES.index(mode)
    pad_enum = GRID_SAMPLE_PADDING_MODES.index(padding_mode)

    if input.dim() == 4:
        return _C.grid_sampler_2d(input, grid, mode_enum, pad_enum, align_corners)
    if input.dim() == 5:
        if mode_enum == 2:
            raise ValueError("nn.functional.grid_sample(): bicubic only supports 4D input")
        return _C.grid_sampler_3d(input, grid, mode_enum, pad_enum, align_corners)
    raise ValueError(f"nn.functional.grid_sample(): expected 4D or 5D input, got {input.dim()}D")

def _grid_sampler_3d(input, grid, interpolation_mode, padding_mode, align_corners):
    if isinstance(padding_mode, str):
        padding_mode = GRID_SAMPLE_PADDING_MODES.index(padding_mode)
    N, C, D_in, H_in, W_in = input.shape
    D_out, H_out, W_out = grid.shape[1], grid.shape[2], grid.shape[3]
    x = _gs_unnormalize(grid[..., 0], W_in, align_corners)
    y = _gs_unnormalize(grid[..., 1], H_in, align_corners)
    z = _gs_unnormalize(grid[..., 2], D_in, align_corners)
    if padding_mode != 0:
        x = _gs_adjust(x, W_in, padding_mode, align_corners)
        y = _gs_adjust(y, H_in, padding_mode, align_corners)
        z = _gs_adjust(z, D_in, padding_mode, align_corners)

    def gather3(xi, yi, zi, ib):
        xi_c = xi.clamp(0, W_in - 1)
        yi_c = yi.clamp(0, H_in - 1)
        zi_c = zi.clamp(0, D_in - 1)
        pos = (zi_c * H_in + yi_c) * W_in + xi_c
        if ib is not None:
            pos = tensorplay.where(ib, pos, pos * 0)
        i64 = DType.int64
        dev = input.device
        bidx = tensorplay.arange(N, dtype=i64, device=dev).view([N, 1, 1, 1, 1]) * (C * D_in * H_in * W_in)
        cidx = tensorplay.arange(C, dtype=i64, device=dev).view([1, C, 1, 1, 1]) * (D_in * H_in * W_in)
        gid = bidx + cidx + pos.unsqueeze(1)
        vals = tensorplay.embedding(input.contiguous().reshape(-1), gid.reshape(-1)).reshape(
            [N, C, D_out, H_out, W_out])
        if ib is not None:
            vals = vals * ib.unsqueeze(1).expand([N, C, D_out, H_out, W_out]).to(vals.dtype)
        return vals

    def bounds3(cx, cy, cz):
        m = tensorplay.logical_and(cx >= 0, cx < W_in)
        m = tensorplay.logical_and(m, tensorplay.logical_and(cy >= 0, cy < H_in))
        m = tensorplay.logical_and(m, tensorplay.logical_and(cz >= 0, cz < D_in))
        return m

    xf = tensorplay.floor(x)
    yf = tensorplay.floor(y)
    zf = tensorplay.floor(z)
    x0, y0, z0 = xf.to(DType.int64), yf.to(DType.int64), zf.to(DType.int64)
    wx, wy, wz = x - xf, y - yf, z - zf

    if interpolation_mode == 1:  # nearest
        xi = tensorplay.round(x).to(DType.int64)
        yi = tensorplay.round(y).to(DType.int64)
        zi = tensorplay.round(z).to(DType.int64)
        ib = bounds3(xi, yi, zi) if padding_mode == 0 else None
        vals = gather3(xi, yi, zi, ib)
        if x.requires_grad:
            # d/dgrid; bridge one so grid.grad stays defined.
            vals = vals + (x * 0).sum()
        return vals

    out = None
    for dz in (0, 1):
        for dy in (0, 1):
            for dx in (0, 1):
                cx, cy, cz = x0 + dx, y0 + dy, z0 + dz
                ib = bounds3(cx, cy, cz) if padding_mode == 0 else None
                v = gather3(cx, cy, cz, ib)
                wgt = (
                    ((wx if dx else 1 - wx) * (wy if dy else 1 - wy) * (wz if dz else 1 - wz))
                    .unsqueeze(1).expand([N, C, D_out, H_out, W_out]).to(v.dtype))
                term = v * wgt
                out = term if out is None else out + term
    return out

# -----------------------------------------------------------------------------
# Distances, embedding bags, SDPA and remaining public surface (F-alignment).
# -----------------------------------------------------------------------------

def _is_float_dtype(dt):
    return dt in (DType.float16, DType.bfloat16, DType.float32, DType.float64)

def _vector_norm(vec, p, keepdim=False):
    """Norm over the last dim of a broadcast difference, matching
"""
    dim = vec.dim() - 1
    if p == float("inf"):
        return _C.max(vec.abs(), dim=dim, keepdim=keepdim)[0]
    if p == -float("inf"):
        return -_C.max((-vec.abs()), dim=dim, keepdim=keepdim)[0]
    if p == 0:
        return vec.ne(0).to(DType.float32).sum(dim=[dim], keepdim=keepdim)
    return _C.norm(vec, [dim], float(p), keepdim)

[docs]
def pairwise_distance(x1: Tensor, x2: Tensor, p: float = 2.0, eps: float = 1e-6, keepdim: bool = False) -> Tensor:
    r"""Computes the pairwise distance between input vectors.

    dimension.
    """
    return _vector_norm(x1 - x2 + eps, p, keepdim)

[docs]
def pdist(input: Tensor, p: float = 2.0) -> Tensor:
    r"""Computes the pairwise distance between rows of :attr:`input`.

    Returns the flattened upper triangle of the ``N x N`` distance matrix —
    """
    if input.dim() != 2:
        raise RuntimeError(f"pdist expects a 2D input, got {input.dim()}D")
    n = input.size(0)
    diff = input.unsqueeze(1) - input.unsqueeze(0)      # (N, N, D)
    dist = _vector_norm(diff, p, keepdim=False)         # (N, N)
    ri = tensorplay.arange(n, dtype=DType.int64, device=input.device)
    mask = ri.view(-1, 1) < ri.view(1, -1)
    return tensorplay.masked_select(dist, mask)

def _no_grad_embedding_renorm_(weight: Tensor, input, max_norm: float, norm_type: float) -> Tensor:
    """Renormalize referenced embedding rows in-place without recording gradients."""
    with tensorplay.no_grad():
        _C.embedding_renorm_(weight, input, float(max_norm), float(norm_type))
    return weight

[docs]
def embedding_bag(
    input: Tensor,
    weight: Tensor,
    offsets=None,
    max_norm=None,
    norm_type: float = 2,
    scale_grad_by_freq: bool = False,
    mode: str = "mean",
    sparse: bool = False,
    per_sample_weights=None,
    include_last_offset: bool = False,
    padding_idx=None,
) -> Tensor:
    r"""Compute sums, means or maxes of ``bags`` of embeddings.

    1-D inputs with :attr:`offsets` (incl. ``include_last_offset``), fixed
    length 2-D inputs, ``per_sample_weights`` (sum mode), ``padding_idx``
    exclusion and ``max_norm`` renormalization.
    See :class:`tensorplay.nn.EmbeddingBag` for details.
    """
    captured = _capture_call(
        embedding_bag,
        (input, weight, offsets, max_norm, norm_type, scale_grad_by_freq, mode,
         sparse, per_sample_weights, include_last_offset, padding_idx),
        {},
    )
    if captured is not None:
        return captured
    # Backward compatibility with the old (weight, input) argument order.
    if weight.dtype == DType.int64 and _is_float_dtype(input.dtype):
        warnings.warn(
            "Argument order of nn.functional.embedding_bag was changed. "
            "Usage `embedding_bag(weight, input, ...)` is deprecated, "
            "and should now be `embedding_bag(input, weight, ...)`.",
            stacklevel=2,
        )
        weight, input = input, weight

    if per_sample_weights is not None and tuple(input.shape) != tuple(per_sample_weights.shape):
        raise ValueError(
            f"embedding_bag: If per_sample_weights ({per_sample_weights.shape}) is not None, "
            f"then it must have the same shape as the input ({input.shape})"
        )
    if weight.dim() != 2:
        raise ValueError(f"weight has to be a 2D Tensor, but got Tensor of dimension {weight.dim()}")

    if mode == "sum":
        mode_enum = 0
    elif mode == "mean":
        mode_enum = 1
    elif mode == "max":
        mode_enum = 2
        if scale_grad_by_freq:
            raise ValueError("max mode does not support scaling the gradient by the frequency")
        if sparse:
            raise ValueError("max mode does not support sparse weights")
    else:
        raise ValueError("mode has to be one of sum, mean or max")

    if max_norm is not None:
        _no_grad_embedding_renorm_(weight, input, max_norm, norm_type)

    if per_sample_weights is not None and mode != "sum":
        raise NotImplementedError(
            "embedding_bag: per_sample_weights was not None. "
            "per_sample_weights is only supported for mode='sum' "
            f"(got mode='{mode}').")

    if padding_idx is not None:
        padding_idx = int(padding_idx)
        if padding_idx >= weight.size(0) or padding_idx < -weight.size(0):
            raise ValueError(
                f"padding_idx must be within the number of embeddings ({weight.size(0)}), "
                f"got {padding_idx}")
        if padding_idx < 0:
            padding_idx += weight.size(0)
    else:
        padding_idx = -1

    i64 = DType.int64
    dev = input.device
    numel = input.numel()
    flat = input.to(i64).reshape(-1)

    include_last_offset = bool(include_last_offset)
    if input.dim() == 2:
        if offsets is not None:
            raise ValueError(
                "if input is 2D, then offsets has to be None"
                ", as input is treated is a mini-batch of"
                " fixed length sequences.")
        seq_len = int(input.size(1))
        # Fixed-length bags: one offset per row.  A zero-width row still owns a
        # bag, which arange cannot express with a zero step.
        if seq_len == 0:
            offs = tensorplay.zeros([input.size(0)], dtype=i64, device=dev)
        else:
            offs = tensorplay.arange(0, numel, seq_len, dtype=i64, device=dev)
        include_last_offset = False
    elif input.dim() == 1:
        if offsets is None:
            raise ValueError("offsets has to be a 1D Tensor but got None")
        if offsets.dim() != 1:
            raise ValueError("offsets has to be a 1D Tensor")
        offs = offsets.to(i64)
    else:
        raise ValueError(
            f"input has to be 1D or 2D Tensor, but got Tensor of dimension {input.dim()}")

    if per_sample_weights is not None:
        per_sample_weights = per_sample_weights.reshape(-1).to(weight.dtype)

    return _C._embedding_bag(
        weight, flat, offs, scale_grad_by_freq, mode_enum, sparse,
        per_sample_weights, include_last_offset, padding_idx)[0]

# The widest head the fused attention tiles are shaped for.  A wider head has
# no tile to land in and stays on the composed reference.
_FUSED_MAX_HEAD_DIM = 128

# Element types the fused cpu attention kernel computes in.
_CPU_FUSED_DTYPES = frozenset(
    (tensorplay.float16, tensorplay.bfloat16, tensorplay.float32, tensorplay.float64))

_AUTOCAST_ELIGIBLE = frozenset((tensorplay.float32, tensorplay.float16, tensorplay.bfloat16))

def _attention_dtype(t):
    # The element type the kernel receives: both attention entries are cast
    # to the autocast type under autocast, so a query left in single precision
    # by a rotary embedding meets its half-precision value there anyway.
    if (
        t.dtype in _AUTOCAST_ELIGIBLE
        and t.device.type == "cuda"
        and tensorplay.is_autocast_enabled("cuda")
    ):
        return tensorplay.get_autocast_dtype("cuda")
    return t.dtype

def _plain_scaled_dot_product_attention(query, key, value, is_causal):
    # The fused entry point takes a self-attention call with the default
    # normaliser: one head count, one token count, and a head width inside the
    # tiled range.  It is also the entry the matching fused backward hangs off,
    # so a training call routed here keeps the fast backward rather than
    # falling back to the composed one.  The token counts may differ -- the
    # schedule reads them independently, like the reference -- and everything
    # else (grouped heads, an explicit normaliser) goes to the public call.
    if (
        query.device.type == "cuda"
        and query.dim() == 4
        and key.dim() == 4
        and value.dim() == 4
        and 0 < query.size(-1) <= _FUSED_MAX_HEAD_DIM
        and _attention_dtype(key) == _attention_dtype(query) == _attention_dtype(value)
        and key.size(-1) == query.size(-1) == value.size(-1)
        and key.size(0) == query.size(0) == value.size(0)
        and key.size(1) == query.size(1) == value.size(1)
    ):
        output, _ = _C._scaled_dot_product_attention_with_lse(
            query, key, value, is_causal=is_causal)
        return output
    # The cpu counterpart is the blocked fused kernel.  It takes any pair of
    # token counts (its causal bound is the top-left one the composed path
    # uses), so only the element type, the head layout and a non-empty
    # sequence are asked for; an empty call is answered by the public entry.
    if (
        query.device.type == "cpu"
        and query.dim() == 4
        and key.dim() == 4
        and value.dim() == 4
        and query.dtype in _CPU_FUSED_DTYPES
        and key.dtype == query.dtype == value.dtype
        and key.size(-1) == query.size(-1) == value.size(-1)
        and key.size(0) == query.size(0) == value.size(0)
        and key.size(1) == query.size(1) == value.size(1)
        and key.size(-2) == value.size(-2)
        and query.numel() != 0
        and key.numel() != 0
    ):
        output, _ = _C._scaled_dot_product_flash_attention_for_cpu(
            query, key, value, 0.0, is_causal, attn_mask=None, scale=None)
        return output
    return tensorplay.scaled_dot_product_attention(
        query, key, value, is_causal=is_causal)

[docs]
def scaled_dot_product_attention(
    query: Tensor,
    key: Tensor,
    value: Tensor,
    attn_mask=None,
    dropout_p: float = 0.0,
    is_causal: bool = False,
    scale=None,
    enable_gqa: bool = False,
    backend: Optional[str] = None,
) -> Tensor:
    r"""scaled_dot_product_attention(query, key, value, attn_mask=None,
    dropout_p=0.0, is_causal=False, scale=None, enable_gqa=False,
    backend=None) -> Tensor

    Computes scaled dot product attention on query, key and value. Routes to
    reference:

    .. math::
        \text{Attention}(Q, K, V) = \text{softmax}(\frac{Q K^T}{\sqrt{E}}) V

    Args:
        backend (str, optional): ``'flash'`` | ``'mem_efficient'``,
            ``'math'``, or ``None`` to pick automatically. ``'flash'``
            selects the fused flash-attention kernel, ``'math'`` forces the
            composed reference path. When ``None``, the routing candidate
            order is governed by :func:`tensorplay.nn.attention.sdpa_kernel`.
    """
    from tensorplay.nn import attention as _sdpa_attention
    from tensorplay.overrides import has_tensorplay_function

    # The plain fused case -- no mask, no drop, no normaliser, no grouping --
    # is the call a decode or training loop makes millions of times.  The
    # routing below re-walks the hook protocol and rebuilds the parameter
    # object for it; a call that carries none of the extras and whose
    # arguments carry no dispatch hook lands on the same fused kernel, so it
    # is answered before any of that is built.
    if (
        attn_mask is None
        and dropout_p == 0.0
        and scale is None
        and not enable_gqa
        and backend is None
        and not has_tensorplay_function((query, key, value))
    ):
        return _plain_scaled_dot_product_attention(query, key, value, is_causal)

    # Attention-bias subclasses intercept the call through the Python
    # function-hook protocol before any backend routing happens.
    if has_tensorplay_function((query, key, value, attn_mask)):
        from tensorplay.overrides import handle_tensorplay_function

        return handle_tensorplay_function(
            scaled_dot_product_attention,
            (query, key, value, attn_mask),
            query,
            key,
            value,
            attn_mask=attn_mask,
            dropout_p=dropout_p,
            is_causal=is_causal,
            scale=scale,
            enable_gqa=enable_gqa,
        )

    if attn_mask is not None and is_causal:
        raise AssertionError("Explicit attn_mask should not be set when is_causal=True")
    if query.dim() < 2 or key.dim() < 2 or value.dim() < 2:
        raise ValueError(
            "scaled_dot_product_attention: query, key and value must be "
            "at least 2-D")
    if key.shape[-1] != query.shape[-1] or value.shape[-1] != query.shape[-1]:
        raise ValueError(
            "scaled_dot_product_attention: query, key and value must share "
            "the last dimension")
    if dropout_p < 0.0 or dropout_p > 1.0:
        raise ValueError(
            f"scaled_dot_product_attention: dropout probability must be in "
            f"[0, 1], got {dropout_p}")
    if backend not in (None, "math", "flash", "mem_efficient"):
        raise ValueError(
            f"scaled_dot_product_attention: unknown backend '{backend}'; "
            "expected 'flash', 'mem_efficient', 'math' or None")

    allowed = _sdpa_attention._cur_sdpa_kernel_backends(with_priority=True)
    full_set = {
        _sdpa_attention.SDPBackend.MATH,
        _sdpa_attention.SDPBackend.FLASH_ATTENTION,
        _sdpa_attention.SDPBackend.EFFICIENT_ATTENTION,
    }
    params = _sdpa_attention.SDPAParams(
        query=query, key=key, value=value, attn_mask=attn_mask,
        dropout=dropout_p, is_causal=is_causal, enable_gqa=enable_gqa,
    )

    if backend is not None:
        _name_to_backend = {
            "math": _sdpa_attention.SDPBackend.MATH,
            "flash": _sdpa_attention.SDPBackend.FLASH_ATTENTION,
            "mem_efficient": _sdpa_attention.SDPBackend.EFFICIENT_ATTENTION,
        }
        requested = _name_to_backend[backend]
        if allowed != full_set and requested not in allowed:
            raise RuntimeError(
                f"scaled_dot_product_attention: backend '{backend}' is "
                f"disabled by the active sdpa_kernel context")
        if backend == "flash":
            plain_case = (
                scale is None
                and attn_mask is None
                and dropout_p == 0.0
                and not enable_gqa
            )
            if plain_case:
                return _plain_scaled_dot_product_attention(
                    query, key, value, is_causal)
            if allowed != full_set and _sdpa_attention.SDPBackend.MATH not in allowed:
                raise RuntimeError(
                    "scaled_dot_product_attention: the requested flash backend "
                    "does not support the supplied mask, scale, or dropout "
                    "configuration")
        if backend == "mem_efficient":
            raise NotImplementedError(
                "scaled_dot_product_attention: the mem_efficient backend requires "
                "native memory-efficient attention kernels, which are not yet "
                "available in this build.")
    else:
        _plain_case = (
            scale is None
            and attn_mask is None
            and dropout_p == 0.0
            and not enable_gqa
        )
        for candidate in allowed:
            if candidate == _sdpa_attention.SDPBackend.FLASH_ATTENTION:
                if _plain_case and _sdpa_attention.can_use_flash_attention(params):
                    # impl=None lets the fused-kernel selection pick the
                    # fastest eligible kernel for the shape.
                    return _plain_scaled_dot_product_attention(
                        query, key, value, is_causal)
                _sdpa_attention._raise_kernel_warnings(params)
            elif candidate == _sdpa_attention.SDPBackend.MATH:
                break
        else:
            raise RuntimeError(
                "No available kernel for scaled_dot_product_attention; "
                "the active sdpa_kernel context allows none of the "
                "installed backends for these inputs")

    # Everything the fused entry point does not serve -- an explicit scale, an
    # explicit mask, grouped heads, a shape the fused kernels cannot express --
    # is answered by the public call.  It carries the derivative that gives
    # this call its gradients, and it reaches the same composed reference from
    # the inside, so a differentiated call and an undifferentiated one are
    # answered by the same computation.  Calling the composed reference
    # directly from here instead would answer an undifferentiated question: no
    # derivative is registered for that entry point, so the result would come
    # back detached and the caller's backward pass would have nothing to
    # differentiate.
    return tensorplay.scaled_dot_product_attention(
        query, key, value, attn_mask=attn_mask,
        dropout_p=dropout_p, is_causal=is_causal, scale=scale,
        enable_gqa=enable_gqa)

[docs]
def linear_cross_entropy(
    input: Tensor,
    linear_weight: Tensor,
    target: Tensor,
    *,
    linear_bias=None,
    weight=None,
    reduction: str = "mean",
    ignore_index=None,
    label_smoothing: float = 0.0,
    options=None,
) -> Tensor:
    r"""Compute cross entropy between ``input``, transformed linearly, and
    target.

    Equivalent to ``cross_entropy(linear(input, linear_weight), target,
    **kwargs)`` (reference path; chunked/fused options are ignored).
    """
    if options is not None:
        warnings.warn(
            "linear_cross_entropy: ``options`` ignored; reference path used.",
            stacklevel=2,
        )
    logits = linear(input, linear_weight, linear_bias)
    ig = -100 if ignore_index is None else ignore_index
    return cross_entropy(logits, target, weight=weight, ignore_index=ig,
                         reduction=reduction, label_smoothing=label_smoothing)

[docs]
def grouped_mm(input, mat2, offs):
    return _C.grouped_mm(input, mat2, offs)

[docs]
def scaled_grouped_mm(*args, **kwargs):
    raise NotImplementedError(
        "scaled_grouped_mm requires native FP8 grouped GEMM kernels, which "
        "are not yet available in this build.")

[docs]
def scaled_mm(*args, **kwargs):
    raise NotImplementedError(
        "scaled_mm requires native FP8 GEMM kernels, which are not yet "
        "available in this build.")

# -----------------------------------------------------------------------------
# -----------------------------------------------------------------------------

def _get_softmax_dim(name: str, ndim: int, stacklevel: int = 3) -> int:
    if ndim == 0 or ndim == 1 or ndim == 3:
        ret = 0
    else:
        ret = 1
    return ret

def _mha_shape_check(query, key, value, key_padding_mask, attn_mask, num_heads):
    if query.dim() == 3:
        bsz, tgt_len, embed_dim_to_check = query.shape
        assert query.shape == (bsz, tgt_len, embed_dim_to_check)
        assert key.shape == value.shape
        bsz, src_len, _ = key.shape
        if key_padding_mask is not None:
            assert key_padding_mask.shape == (bsz, src_len)
            assert key_padding_mask.dtype == DType.bool
        if attn_mask is not None:
            if attn_mask.dim() == 2:
                correct_2d_size = (tgt_len, src_len)
                assert attn_mask.shape == correct_2d_size
            elif attn_mask.dim() == 3:
                correct_3d_size = (bsz * num_heads, tgt_len, src_len)
                assert attn_mask.shape == correct_3d_size
    elif query.dim() == 2:
        assert key.dim() == 2
        assert key.shape == value.shape
        src_len, _ = key.shape
        if key_padding_mask is not None:
            assert key_padding_mask.shape == (src_len,)
            assert key_padding_mask.dtype == DType.bool
        if attn_mask is not None:
            correct_2d_size = (src_len, src_len)
            if attn_mask.dim() == 2:
                assert attn_mask.shape == correct_2d_size
    else:
        raise AssertionError(
            f"query has to be 2d or 3d, but got {query.dim()}d")

def _in_projection(
    q,
    k,
    v,
    w_q=None,
    w_k=None,
    w_v=None,
    b_q=None,
    b_k=None,
    b_v=None,
):
    """
    projections with shape constraints ensuring embedding uniformity."""
    Eq, Ek, Ev = q.size(-1), k.size(-1), v.size(-1)
    if w_q.shape != (Eq, Eq):
        raise AssertionError(f"expecting query weights shape of {(Eq, Eq)}, but got {tuple(w_q.shape)}")
    if w_k.shape != (Eq, Ek):
        raise AssertionError(f"expecting key weights shape of {(Eq, Ek)}, but got {tuple(w_k.shape)}")
    if w_v.shape != (Eq, Ev):
        raise AssertionError(f"expecting value weights shape of {(Eq, Ev)}, but got {tuple(w_v.shape)}")
    if b_q is not None and b_q.shape != (Eq,):
        raise AssertionError(f"expecting query bias shape of {(Eq,)}, but got {tuple(b_q.shape)}")
    if b_k is not None and b_k.shape != (Eq,):
        raise AssertionError(f"expecting key bias shape of {(Eq,)}, but got {tuple(b_k.shape)}")
    if b_v is not None and b_v.shape != (Eq,):
        raise AssertionError(f"expecting value bias shape of {(Eq,)}, but got {tuple(b_v.shape)}")
    return linear(q, w_q, b_q), linear(k, w_k, b_k), linear(v, w_v, b_v)

def _in_projection_packed(q, k, v, w, b=None):
    E = q.size(-1)
    if k is v:
        if q is k:
            # self-attention: one packed projection. Reshape (not chunk) so an
            # unflatten(-1, (3, E)) validation.
            proj = linear(q, w, b)
            p = proj.reshape(tuple(proj.shape[:-1]) + (3, E))
            return tuple(p.select(-2, i).contiguous() for i in range(3))
        # encoder-decoder attention: packed kv + separate q (split is exact).
        if w.size(0) != 3 * E:
            raise RuntimeError(
                f"in_projection_packed: expected packed weight with 3*E={3 * E} "
                f"rows, got {w.size(0)}")
        w_q = w.narrow(0, 0, E)
        w_kv = w.narrow(0, E, E * 2)
        if b is None:
            b_q = b_kv = None
        else:
            b_q = b.narrow(0, 0, E)
            b_kv = b.narrow(0, E, E * 2)
        q_proj = linear(q, w_q, b_q)
        kv_proj = linear(k, w_kv, b_kv)
        p_k, p_v = tensorplay.chunk(kv_proj, 2, dim=-1)
        return q_proj, p_k.contiguous(), p_v.contiguous()
    # (3, -1), which requires the row count to be divisible by three).
    if w.size(0) % 3 != 0:
        raise RuntimeError(
            f"in_projection_packed: packed weight rows ({w.size(0)}) must be "
            f"divisible by 3")
    w_chunks = tensorplay.chunk(w, 3, dim=0)
    if b is None:
        b_q = b_k = b_v = None
    else:
        b_chunks = tensorplay.chunk(b, 3, dim=0)
        b_q, b_k, b_v = b_chunks[0], b_chunks[1], b_chunks[2]
    return (
        linear(q, w_chunks[0], b_q),
        linear(k, w_chunks[1], b_k),
        linear(v, w_chunks[2], b_v),
    )

# -----------------------------------------------------------------------------
# Pooling family (F-alignment).  Native kernels exist for 2-D pooling only;
# 3-D pools decompose into 2-D + 1-D stages (window placement factors per
# Index-returning variants reuse the native/composed values and recover
# -----------------------------------------------------------------------------

def _pool_out_size(in_size, k, s, p, d, ceil_mode):
    eff_k = d * (k - 1) + 1
    if ceil_mode:
        out = int(math.ceil((in_size + 2 * p - eff_k) / s)) + 1
        if (out - 1) * s >= in_size + p:
            out -= 1
    else:
        out = int(math.floor((in_size + 2 * p - eff_k) / s)) + 1
    return max(out, 0)

def _max_pool2d_indices(x4, kernel_size, stride, padding, dilation, oH=None, oW=None):
    """Per-plane linear indices of a 2-D max pool over ``(N, C, H, W)``.

    ``row * W + col``; first occurrence wins ties.
    """
    with tensorplay.no_grad():
        N, C, H, W = x4.shape
        kh, kw = kernel_size
        sh, sw = stride
        ph, pw = padding
        dh, dw = dilation
        dev = x4.device
        i64 = DType.int64
        if oH is None:
            oH = _pool_out_size(H, kh, sh, ph, dh, False)
        if oW is None:
            oW = _pool_out_size(W, kw, sw, pw, dw, False)

        def _rng(n):
            return tensorplay.arange(n, dtype=i64, device=dev)

        Rm = (_rng(oH) * sh).view(-1, 1, 1, 1) + (_rng(kh) * dh).view(1, 1, -1, 1) - ph
        Cm = (_rng(oW) * sw).view(1, -1, 1, 1) + (_rng(kw) * dw).view(1, 1, 1, -1) - pw
        gR = (Rm + Cm * 0).reshape(1, oH * oW, kh * kw)
        gC = (Cm + Rm * 0).reshape(1, oH * oW, kh * kw)
        valid = tensorplay.logical_and(_band(gR, 0, H), _band(gC, 0, W))

        pos = gR.clamp(0, H - 1) * W + gC.clamp(0, W - 1)
        pos = tensorplay.where(valid, pos, pos * 0)

        P = N * C
        M, K = oH * oW, kh * kw
        base = (_rng(P) * (H * W)).view(P, 1)
        gid = base + pos.reshape(1, M * K)
        vals = tensorplay.embedding(x4.contiguous().reshape(-1), gid).view(P, M, K)
        vals = tensorplay.where(valid, vals, tensorplay.full_like(vals, float("-inf")))
        # argmax over the kernel axis (last of (P, M, K)): one-hot select of
        # the winning offset per output position.
        am = _C.argmax(vals, 2, False)

        kar = _rng(K).view(1, 1, K)
        sel = am.unsqueeze(-1).eq(kar)
        pick_r = tensorplay.where(sel, gR, gR * 0).sum(-1)
        pick_c = tensorplay.where(sel, gC, gC * 0).sum(-1)
        ok = tensorplay.logical_and(_band(pick_r, 0, H), _band(pick_c, 0, W))
        idx = tensorplay.where(ok, pick_r * W + pick_c, pick_r * 0)
        return idx.view(N, C, oH, oW)

[docs]
def max_pool2d_with_indices(
    input: Tensor,
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode: bool = False,
    return_indices: bool = True,
):
    r"""Applies a 2D max pooling over an input composed of several input
    planes, returning ``(output, indices)``.

    See :class:`~tensorplay.nn.MaxPool2d` for details.
    """
    kernel_size = _pair(kernel_size)
    stride = kernel_size if stride is None else _pair(stride)
    padding = _pair(padding)
    dilation = _pair(dilation)
    # int64 indices into each (n, c) input plane; autograd flows through
    # max_pool2d_with_indices_backward.
    return _C.max_pool2d_with_indices(input, list(kernel_size), list(stride),
                                      list(padding), list(dilation), ceil_mode)

[docs]
def max_pool1d_with_indices(
    input: Tensor,
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode: bool = False,
    return_indices: bool = True,
):
    r"""Applies a 1D max pooling over an input signal, returning
    ``(output, indices)``.

    See :class:`~tensorplay.nn.MaxPool1d` for details.
    """
    unbatched = input.dim() == 2
    x = input.unsqueeze(0) if unbatched else input
    k = _single(kernel_size)[0]
    s = k if stride is None else _single(stride)[0]
    p = _single(padding)[0]
    d = _single(dilation)[0]
    values, indices = max_pool2d_with_indices(
        x.unsqueeze(3), (k, 1), (s, 1), (p, 0), (d, 1), ceil_mode
    )
    values = values.squeeze(3)
    indices = indices.squeeze(3)
    if unbatched:
        return values.squeeze(0), indices.squeeze(0)
    return values, indices

[docs]
def max_pool3d(
    input: Tensor,
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode: bool = False,
    return_indices: bool = False,
) -> Tensor:
    r"""max_pool3d(input, kernel_size, stride=None, padding=0, dilation=1,
    ceil_mode=False, return_indices=False) -> Tensor

    Applies a 3D max pooling over an input signal composed of several input
    planes. Input shape ``(N, C, D, H, W)`` or unbatched ``(C, D, H, W)``.
    """
    captured = _capture_call(max_pool3d, (input, kernel_size, stride, padding, dilation, ceil_mode, return_indices), {})
    if captured is not None:
        return captured
    if return_indices:
        return max_pool3d_with_indices(
            input, kernel_size, stride=stride, padding=padding,
            dilation=dilation, ceil_mode=ceil_mode)
    kd, kh, kw = _triple(kernel_size)
    sd, sh, sw = _triple(stride) if stride is not None else (kd, kh, kw)
    pd_, ph, pw = _triple(padding)
    dd, dh, dw = _triple(dilation)
    # max_pool3d_backward.
    return _C.max_pool3d(input, [kd, kh, kw], [sd, sh, sw], [pd_, ph, pw],
                         [dd, dh, dw], ceil_mode)

[docs]
def max_pool3d_with_indices(
    input: Tensor,
    kernel_size,
    stride=None,
    padding=0,
    dilation=1,
    ceil_mode: bool = False,
    return_indices: bool = True,
):
    r"""Applies a 3D max pooling over an input signal, returning
    ``(output, indices)``.

    See :class:`~tensorplay.nn.MaxPool3d` for details.
    """
    kd, kh, kw = _triple(kernel_size)
    sd, sh, sw = _triple(stride) if stride is not None else (kd, kh, kw)
    pd_, ph, pw = _triple(padding)
    dd, dh, dw = _triple(dilation)
    # Native kernel returns values plus int64 indices into each (n, c) input
    # (D, H, W) volume; autograd flows through max_pool3d_with_indices_backward.
    return _C.max_pool3d_with_indices(input, [kd, kh, kw], [sd, sh, sw],
                                      [pd_, ph, pw], [dd, dh, dw], ceil_mode)

[docs]
def avg_pool3d(
    input: Tensor,
    kernel_size,
    stride=None,
    padding=0,
    ceil_mode: bool = False,
    count_include_pad: bool = True,
    divisor_override=None,
) -> Tensor:
    r"""avg_pool3d(input, kernel_size, stride=None, padding=0, ceil_mode=False,
    count_include_pad=True, divisor_override=None) -> Tensor

    Applies a 3D average pooling over an input signal composed of several
    input planes. Input shape ``(N, C, D, H, W)`` or unbatched ``(C, D, H, W)``.
    """
    captured = _capture_call(avg_pool3d, (input, kernel_size, stride, padding, ceil_mode, count_include_pad, divisor_override), {})
    if captured is not None:
        return captured
    unbatched = input.dim() == 4
    x = input.unsqueeze(0) if unbatched else input
    kd, kh, kw = _triple(kernel_size)
    if stride is None:
        sd, sh, sw = kd, kh, kw
    else:
        sd, sh, sw = _triple(stride)
    pd_, ph, pw = _triple(padding)
    output = _C.avg_pool3d(x, (kd, kh, kw), (sd, sh, sw), (pd_, ph, pw),
                            ceil_mode, count_include_pad, divisor_override)
    return output.squeeze(0) if unbatched else output

[docs]
def adaptive_avg_pool3d(input: Tensor, output_size) -> Tensor:
    r"""Apply a 3D adaptive average pooling over an input signal."""
    unbatched = input.dim() == 4
    x = input.unsqueeze(0) if unbatched else input
    od, oh, ow = _triple(output_size)
    return _C.adaptive_avg_pool3d(x, (od, oh, ow)).squeeze(0) if unbatched \
        else _C.adaptive_avg_pool3d(x, (od, oh, ow))

[docs]
def lp_pool3d(
    input: Tensor,
    norm_type,
    kernel_size,
    stride=None,
    ceil_mode: bool = False,
) -> Tensor:
    r"""Apply a 3D power-average pooling over an input signal.

    See :class:`~tensorplay.nn.LPPool3d` for details.
    """
    kd, kw, kh = _triple(kernel_size)
    if isinstance(norm_type, (int, float)):
        if norm_type == 0:
            raise ValueError(f"norm_type must be a non-zero value, but got {norm_type}")
        if norm_type == float("inf"):
            return max_pool3d(input.abs(), kernel_size, stride, 0, 1, ceil_mode)
        if norm_type == -float("inf"):
            return -max_pool3d((-input.abs()), kernel_size, stride, 0, 1, ceil_mode)

    if stride is not None:
        out = avg_pool3d(input.pow(norm_type), kernel_size, stride, 0, ceil_mode)
    else:
        out = avg_pool3d(input.pow(norm_type), kernel_size, padding=0, ceil_mode=ceil_mode)

    return (tensorplay.sign(out) * relu(tensorplay.abs(out))).mul(kd * kw * kh).pow(1.0 / norm_type)

[docs]
def adaptive_max_pool2d_with_indices(input: Tensor, output_size, return_indices: bool = True):
    r"""Applies a 2D adaptive max pooling over an input signal composed of
    several input planes, returning ``(output, indices)``.

    See :class:`~tensorplay.nn.AdaptiveMaxPool2d` for details.
    """
    output_size = list(_pair(output_size))
    unbatched = input.dim() == 3
    x = input.unsqueeze(0) if unbatched else input
    values, indices = _C.adaptive_max_pool2d_with_indices(x, output_size)
    if unbatched:
        return values.squeeze(0), indices.squeeze(0)
    return values, indices

[docs]
def adaptive_max_pool1d_with_indices(input: Tensor, output_size, return_indices: bool = True):
    r"""Applies a 1D adaptive max pooling over an input signal, returning
    ``(output, indices)``.

    See :class:`~tensorplay.nn.AdaptiveMaxPool1d` for details.
    """
    unbatched = input.dim() == 2
    x = input.unsqueeze(0) if unbatched else input
    values, indices = adaptive_max_pool2d_with_indices(x.unsqueeze(3), (output_size, 1))
    values = values.squeeze(3)
    indices = indices.squeeze(3)
    if unbatched:
        return values.squeeze(0), indices.squeeze(0)
    return values, indices

[docs]
def adaptive_max_pool3d(input: Tensor, output_size, return_indices: bool = False):
    r"""adaptive_max_pool3d(input, output_size, return_indices=False)

    Applies a 3D adaptive max pooling over an input signal composed of several
    input planes. Input shape ``(N, C, D, H, W)`` or unbatched ``(C, D, H, W)``.

    See :class:`~tensorplay.nn.AdaptiveMaxPool3d` for details.
    """
    od, oh, ow = _triple(output_size)
    values, indices = _C.adaptive_max_pool3d(input, [od, oh, ow])
    if return_indices:
        return values, indices
    return values

[docs]
def adaptive_max_pool3d_with_indices(input: Tensor, output_size, return_indices: bool = True):
    r"""Applies a 3D adaptive max pooling over an input signal, returning
    ``(output, indices)``.

    See :class:`~tensorplay.nn.AdaptiveMaxPool3d` for details.
    """
    od, oh, ow = _triple(output_size)
    unbatched = input.dim() == 4
    x = input.unsqueeze(0) if unbatched else input
    values, indices = _C.adaptive_max_pool3d(x, [od, oh, ow])
    if unbatched:
        return values.squeeze(0), indices.squeeze(0)
    return values, indices

# -----------------------------------------------------------------------------
# generate_intervals + FractionalMaxPool{2d,3d}.cpp window scan, vectorized
# through an embedding gather so autograd flows to the input).
# -----------------------------------------------------------------------------

def _frac_generate_intervals(sample, in_size, out_size, pool_size):
    """

    ``sample`` is a ``(P,)`` float tensor; returns ``(P, out_size)`` int64
    window start positions.
    """
    dev = sample.device
    last = in_size - pool_size
    P = sample.numel()
    if out_size <= 1:
        return tensorplay.full([P, max(out_size, 0)], last, dtype=DType.int64, device=dev)
    alpha = float(in_size - pool_size) / float(out_size - 1)
    i = tensorplay.arange(out_size - 1, dtype=sample.dtype, device=dev)
    # static_cast<int> truncates toward zero; operands are non-negative here
    seq = ((i.unsqueeze(0) + sample.unsqueeze(1)) * alpha).to(DType.int64)
    seq = seq - (sample * alpha).to(DType.int64).unsqueeze(1)
    tail = tensorplay.full([P, 1], last, dtype=DType.int64, device=dev)
    return tensorplay.cat([seq, tail], dim=1)

def _frac_pool_check(input, _random_samples, ndim_spatial):
    if _random_samples.dim() != 3:
        raise ValueError(f"Expect _random_samples to have 3 dimensions, got {_random_samples.dim()}")
    nbatch = 1 if input.dim() == ndim_spatial + 1 else input.size(0)
    channels = input.size(0) if input.dim() == ndim_spatial + 1 else input.size(1)
    if _random_samples.size(0) < nbatch:
        raise ValueError("Expect _random_samples.size(0) no less then input batch size.")
    if _random_samples.size(1) != channels:
        raise ValueError("Expect _random_samples.size(1) equals to input channel size.")
    if _random_samples.size(2) != ndim_spatial:
        raise ValueError(f"Expect _random_samples.size(2) equals to {ndim_spatial}; got {_random_samples.size(2)}.")

def _frac_windowed_max(x_planes, pos, plane_size):
    """Max + argmax of windows gathered per-plane through an embedding lookup.

    x_planes: ``(P, V)``; pos: ``(P, M, K)`` int64 within-plane positions.
    ``val > maxVal || isnan(val)`` scan.
    """
    P = x_planes.size(0)
    M, K = pos.shape[1], pos.shape[2]
    base = (tensorplay.arange(P, dtype=DType.int64, device=x_planes.device) * plane_size).view(P, 1)
    gid = base + pos.reshape(P, M * K)
    vals = tensorplay.embedding(x_planes.contiguous().reshape(-1), gid).view(P, M, K)
    sub = tensorplay.where(vals.ne(vals), tensorplay.full_like(vals, float("inf")), vals)
    am = _C.argmax(sub, 2, False)  # window-max index along K
    kar = tensorplay.arange(K, dtype=DType.int64, device=x_planes.device).view(1, 1, K)
    sel = am.unsqueeze(-1).eq(kar)
    picked = tensorplay.where(sel, vals, vals * 0).sum(-1)
    idx = tensorplay.where(sel, pos, pos * 0).sum(-1)
    return picked, idx

[docs]
def fractional_max_pool2d_with_indices(
    input: Tensor,
    kernel_size,
    output_size=None,
    output_ratio=None,
    return_indices: bool = True,
    _random_samples=None,
):
    r"""Applies 2D fractional max pooling over an input signal composed of
    several input planes, returning ``(output, indices)``.

    The max-pooling operation is applied in :math:`kH \times kW` regions by a
    stochastic step size determined by the target output size.
    The number of output features is equal to the number of input planes.

    Args:
        kernel_size: the size of the window, ``k`` or ``(kH, kW)``
        output_size: target output size ``oH x oW``
        output_ratio: alternative to output_size, in range (0, 1)
        return_indices: return pooling indices as well
        _random_samples: optional ``(B, C, 2)`` random starts override

    See :class:`~tensorplay.nn.FractionalMaxPool2d` for details.
    """
    if output_size is None and output_ratio is None:
        raise ValueError("fractional_max_pool2d requires specifying either an output_size or an output_ratio")
    if output_size is None:
        _output_ratio = _pair(output_ratio)
        output_size = [int(input.size(-2) * _output_ratio[0]), int(input.size(-1) * _output_ratio[1])]
    output_size = list(_pair(output_size))
    kh, kw = _pair(kernel_size)
    oH, oW = int(output_size[0]), int(output_size[1])

    unbatched = input.dim() == 3
    x = input.unsqueeze(0) if unbatched else input
    B, C, H, W = x.shape
    if oH < 1 or oW < 1 or kh < 1 or kw < 1:
        raise ValueError(
            f"fractional_max_pool2d: kernel_size ({kh}, {kw}) and output_size ({oH}, {oW}) must be positive")
    if oH + kh - 1 > H or oW + kw - 1 > W:
        raise ValueError(
            f"fractional_max_pool2d: output_size ({oH}, {oW}) too large relative to "
            f"input ({H}, {W}) and kernel ({kh}, {kw})")

    if _random_samples is None:
        _random_samples = tensorplay.rand(B, C, 2, dtype=input.dtype, device=input.device)
    _frac_pool_check(x, _random_samples, 2)

    # intervals derive from _random_samples, indices are flat in-plane offsets.
    values, indices = _C.fractional_max_pool2d(x, [kh, kw], [oH, oW], _random_samples)
    if unbatched:
        return values.squeeze(0), indices.squeeze(0)
    return values, indices

def _fractional_max_pool2d(
    input: Tensor,
    kernel_size,
    output_size=None,
    output_ratio=None,
    return_indices: bool = False,
    _random_samples=None,
) -> Tensor:
    return fractional_max_pool2d_with_indices(
        input, kernel_size, output_size=output_size, output_ratio=output_ratio,
        return_indices=return_indices, _random_samples=_random_samples)[0]

[docs]
def fractional_max_pool2d(
    input: Tensor,
    kernel_size,
    output_size=None,
    output_ratio=None,
    return_indices: bool = False,
    _random_samples=None,
):
    r"""Applies 2D fractional max pooling over an input signal.

    If :attr:`return_indices` is ``True``, returns ``(output, indices)``;
    otherwise just the output.
    """
    if return_indices:
        return fractional_max_pool2d_with_indices(
            input, kernel_size, output_size=output_size, output_ratio=output_ratio,
            return_indices=True, _random_samples=_random_samples)
    return _fractional_max_pool2d(
        input, kernel_size, output_size=output_size, output_ratio=output_ratio,
        _random_samples=_random_samples)

[docs]
def fractional_max_pool3d_with_indices(
    input: Tensor,
    kernel_size,
    output_size=None,
    output_ratio=None,
    return_indices: bool = True,
    _random_samples=None,
):
    r"""Applies 3D fractional max pooling over an input signal composed of
    several input planes, returning ``(output, indices)``.

    Each plane consumes three random samples ordered ``(T, H, W)``, matching
    """
    if output_size is None and output_ratio is None:
        raise ValueError("fractional_max_pool3d requires specifying either an output_size or an output_ratio")
    if output_size is None:
        _output_ratio = _triple(output_ratio)
        output_size = [
            int(input.size(-3) * _output_ratio[0]),
            int(input.size(-2) * _output_ratio[1]),
            int(input.size(-1) * _output_ratio[2]),
        ]
    output_size = list(_triple(output_size))
    kt, kh, kw = _triple(kernel_size)
    oT, oH, oW = int(output_size[0]), int(output_size[1]), int(output_size[2])

    unbatched = input.dim() == 4
    x = input.unsqueeze(0) if unbatched else input
    B, C, T, H, W = x.shape
    if min(oT, oH, oW) < 1 or min(kt, kh, kw) < 1:
        raise ValueError("fractional_max_pool3d: kernel_size and output_size must be positive")
    if oT + kt - 1 > T or oH + kh - 1 > H or oW + kw - 1 > W:
        raise ValueError(
            f"fractional_max_pool3d: output_size ({oT}, {oH}, {oW}) too large relative to "
            f"input ({T}, {H}, {W}) and kernel ({kt}, {kh}, {kw})")

    if _random_samples is None:
        _random_samples = tensorplay.rand(B, C, 3, dtype=input.dtype, device=input.device)
    _frac_pool_check(x, _random_samples, 3)

    # samples ordered (T, H, W), indices flat in-plane offsets.
    values, indices = _C.fractional_max_pool3d(x, [kt, kh, kw], [oT, oH, oW], _random_samples)
    if unbatched:
        return values.squeeze(0), indices.squeeze(0)
    return values, indices

def _fractional_max_pool3d(
    input: Tensor,
    kernel_size,
    output_size=None,
    output_ratio=None,
    return_indices: bool = False,
    _random_samples=None,
) -> Tensor:
    return fractional_max_pool3d_with_indices(
        input, kernel_size, output_size=output_size, output_ratio=output_ratio,
        return_indices=return_indices, _random_samples=_random_samples)[0]

[docs]
def fractional_max_pool3d(
    input: Tensor,
    kernel_size,
    output_size=None,
    output_ratio=None,
    return_indices: bool = False,
    _random_samples=None,
):
    r"""Applies 3D fractional max pooling over an input signal.

    If :attr:`return_indices` is ``True``, returns ``(output, indices)``;
    otherwise just the output.
    """
    if return_indices:
        return fractional_max_pool3d_with_indices(
            input, kernel_size, output_size=output_size, output_ratio=output_ratio,
            return_indices=True, _random_samples=_random_samples)
    return _fractional_max_pool3d(
        input, kernel_size, output_size=output_size, output_ratio=output_ratio,
        _random_samples=_random_samples)

# -----------------------------------------------------------------------------
# Max unpooling (partial inverse of max pooling).  Composed as a scatter of
# the pooled values into a zero canvas via the differentiable index_add op.
# -----------------------------------------------------------------------------

def _unpool_output_size(
    input: Tensor,
    kernel_size,
    stride,
    padding,
    output_size,
):
    input_size = input.size()
    n = len(kernel_size)
    default_size = [
        (input_size[-n + d] - 1) * stride[d] + kernel_size[d] - 2 * padding[d]
        for d in range(n)
    ]
    if output_size is None:
        ret = default_size
    else:
        if len(output_size) == n + 2:
            output_size = output_size[2:]
        if len(output_size) != n:
            raise ValueError(
                "output_size should be a sequence containing "
                f"{n} or {n + 2} elements, but it has a length of '{len(output_size)}'"
            )
        ret = list(output_size)
        for d in range(n):
            min_size = default_size[d] - stride[d]
            max_size = default_size[d] + stride[d]
            if not (min_size < ret[d] < max_size):
                raise ValueError(
                    f'invalid output_size "{output_size}" (dim {d} must be between {min_size} and {max_size})'
                )
    for d in range(n):
        if ret[d] < 0:
            raise ValueError(
                "max_unpooling: output_size must contain non-negative spatial "
                f"dimensions, but got output_size[{d}]={ret[d]}"
            )
    return ret

[docs]
def max_unpool1d(
    input: Tensor,
    indices: Tensor,
    kernel_size,
    stride=None,
    padding=0,
    output_size=None,
) -> Tensor:
    r"""Compute a partial inverse of :class:`MaxPool1d`.

    See :class:`~tensorplay.nn.MaxUnpool1d` for details.
    """
    kernel_size = _single(kernel_size)
    _stride = _single(stride) if stride is not None else kernel_size
    padding = _single(padding)
    output_size = _unpool_output_size(input, kernel_size, _stride, padding, output_size)
    return _C.max_unpool2d(
        input.unsqueeze(-1), indices.unsqueeze(-1), list(output_size) + [1]
    ).squeeze(-1)

[docs]
def max_unpool2d(
    input: Tensor,
    indices: Tensor,
    kernel_size,
    stride=None,
    padding=0,
    output_size=None,
) -> Tensor:
    r"""Compute a partial inverse of :class:`MaxPool2d`.

    See :class:`~tensorplay.nn.MaxUnpool2d` for details.
    """
    kernel_size = _pair(kernel_size)
    _stride = _pair(stride) if stride is not None else kernel_size
    padding = _pair(padding)
    output_size = _unpool_output_size(input, kernel_size, _stride, padding, output_size)
    # pooled values into a zero canvas at the flat in-plane int64 indices.
    return _C.max_unpool2d(input, indices, list(output_size))

[docs]
def max_unpool3d(
    input: Tensor,
    indices: Tensor,
    kernel_size,
    stride=None,
    padding=0,
    output_size=None,
) -> Tensor:
    r"""Compute a partial inverse of :class:`MaxPool3d`.

    See :class:`~tensorplay.nn.MaxUnpool3d` for details.
    """
    kernel_size = _triple(kernel_size)
    _stride = _triple(stride) if stride is not None else kernel_size
    padding = _triple(padding)
    output_size = _unpool_output_size(input, kernel_size, _stride, padding, output_size)
    return _C.max_unpool3d(input, indices, list(output_size), list(_stride), list(padding))

class ScalingType(enum.Enum):
    """How a scale is spread across the values it scales.

    A scale is not one thing: one number for a whole tile, one per row, one per
    row of blocks, and one per square block are four different scales, and a
    kernel that reads one of them reads the others wrongly rather than not at
    all.  So the shape a scale has is named, and a kernel is told which name it
    is being handed.

    The names are the shapes' own rather than the algorithms' -- ``BlockWise1x128``
    says one scale per row of blocks with 128 in the other direction, and says
    nothing about how a kernel is to read it.  That is deliberate: a fifth shape
    would be a fifth name here and a fifth branch in every kernel that reads one.
    """

    #: One number for the whole tile.
    TensorWise = "tensorwise"
    #: One number per row of the side being scaled.
    RowWise = "rowwise"
    #: One number per row of blocks, and 128 within a row.
    BlockWise1x128 = "blockwise1x128"
    #: One number per 128 by 128 block.
    BlockWise128x128 = "blockwise128x128"
    #: One number per row of blocks, and 32 within a row.
    #:
    #: The narrow block is what a format with a scale per group of values uses,
    #: where the group's own scale is itself a value of that format rather than a
    #: number.  It is named here so that a call using it is *recognised* as using
    #: it -- which is what lets the call be handed to the operation that can
    #: compute it, rather than to a template that would read the group scale as
    #: if it were the only one.
    BlockWise1x32 = "blockwise1x32"
    #: One number per row of blocks, and 16 within a row.
    BlockWise1x16 = "blockwise1x16"
```