Functional API¶

The functional module provides a collection of stateless functions that perform common operations on tensors. These functions are the functional counterparts to the Layer classes and can be used directly without creating layer objects.

#include "mllm/nn/Functional.hpp"

namespace mllm::nn::functional

Matrix Operations¶

Tensor mllm::nn::functional::matmul(const Tensor &A, const Tensor &B, bool transpose_A = false, bool transpose_B = false, aops::MatMulOpType type = aops::MatMulOpType::kDefault)¶

Perform matrix multiplication of two tensors.

Parameters:

A – First input tensor
B – Second input tensor
transpose_A – Whether to transpose tensor A before multiplication (default: false)
transpose_B – Whether to transpose tensor B before multiplication (default: false)
type – Type of matrix multiplication operation (default: kDefault)

Returns:

Result of matrix multiplication

Shape Operations¶

Tensor mllm::nn::functional::view(const Tensor &x, const std::vector<int32_t> &shape)¶

Reshape a tensor to a new shape.

Parameters:

x – Input tensor
shape – New shape for the tensor

Returns:

Reshaped tensor

std::vector<Tensor> mllm::nn::functional::split(const Tensor &x, int32_t split_size_or_sections, int32_t dim)¶

Split a tensor into chunks along a given dimension.

Parameters:

x – Input tensor
split_size_or_sections – Size of each chunk or list of sizes for each chunk
dim – Dimension along which to split the tensor

Returns:

Vector of split tensors

std::vector<Tensor> mllm::nn::functional::split(const Tensor &x, const std::vector<int32_t> &split_size_or_sections, int32_t dim)¶

Split a tensor into chunks with specified sizes along a given dimension.

Parameters:

x – Input tensor
split_size_or_sections – List of sizes for each chunk
dim – Dimension along which to split the tensor

Returns:

Vector of split tensors

template<int32_t RET_NUM> std::array<Tensor, RET_NUM> mllm::nn::functional::split(const Tensor &x, int32_t split_size_or_sections, int32_t dim)¶

Split a tensor into a fixed number of chunks with same size along a given dimension.

Parameters:

x – Input tensor
split_size_or_sections – Size of each chunk
dim – Dimension along which to split the tensor

Returns:

Array of split tensors with fixed size

template<int32_t RET_NUM> std::array<Tensor, RET_NUM> mllm::nn::functional::split(const Tensor &x, const std::vector<int32_t> &split_size_or_sections, int32_t dim)¶

Split a tensor into a fixed number of chunks with specified sizes along a given dimension.

Parameters:

x – Input tensor
split_size_or_sections – List of sizes for each chunk
dim – Dimension along which to split the tensor

Returns:

Array of split tensors with fixed size

Tensor mllm::nn::functional::concat(const std::vector<Tensor> &ins, int32_t dim)¶

Concatenate a sequence of tensors along a given dimension.

Parameters:

ins – Vector of input tensors to concatenate
dim – Dimension along which to concatenate

Returns:

Concatenated tensor

Tensor mllm::nn::functional::pad(const Tensor &x, const std::vector<int32_t> &pad, aops::PadMode mode = aops::PadMode::kConstant, float value = 0.0f)¶

Pad a tensor along the last N dimensions as specified.

Parameters:

x – Input tensor
pad – Padding sizes ordered from the last dimension to the first, e.g. [last_left, last_right, …, first_left, first_right]
mode – Padding mode (kConstant, kReflect, kReplicate, kCircular). Default: kConstant
value – Constant value used when mode is kConstant. Default: 0.0

Returns:

Padded tensor

Tensor mllm::nn::functional::interpolate(const Tensor &x, const std::vector<int32_t> &size, aops::InterpolateOpMode mode = aops::InterpolateOpMode::kNearest, bool align_corners = false, bool antialias = false)¶

Resize a tensor to the target spatial size.

Parameters:

x – Input tensor (supports 1D/2D/3D spatial resizing depending on mode)
size – Target spatial size (e.g., [H_out, W_out] for 2D)
mode – Interpolation mode (kNearest, kLinear, kBilinear, kBicubic, kTrilinear). Default: kNearest
align_corners – Align corners for linear/bilinear/trilinear interpolation. Default: false :return: Resized tensor

Tensor mllm::nn::functional::interpolate(const Tensor &x, const std::vector<float> &scale_factor, aops::InterpolateOpMode mode = aops::InterpolateOpMode::kNearest, bool align_corners = false)¶

Resize a tensor by scale factors per spatial dimension.

Parameters:

x – Input tensor (supports 1D/2D/3D spatial resizing depending on mode)
scale_factor – Scale factors per spatial dimension (e.g., [sh, sw] for 2D)
mode – Interpolation mode (kNearest, kLinear, kBilinear, kBicubic, kTrilinear). Default: kNearest
align_corners – Align corners for linear/bilinear/trilinear interpolation. Default: false

Returns:

Resized tensor

Attention Operations¶

Tensor mllm::nn::functional::flashAttention2(const Tensor &Q, const Tensor &K, const Tensor &V)¶

Perform FlashAttention-2 operation on query, key, and value tensors.

Parameters:

Q – Query tensor in BSHD format
K – Key tensor in BSHD format
V – Value tensor in BSHD format

Returns:

Output tensor after attention operation

Activation Functions¶

Tensor mllm::nn::functional::softmax(const Tensor &x, int32_t dim)¶

Apply softmax activation function along a given dimension.

Parameters:

x – Input tensor
dim – Dimension along which to apply softmax

Returns:

Tensor with softmax applied

Tensor mllm::nn::functional::log(const Tensor &x)¶

Compute natural logarithm of elements in the tensor.

Parameters:: x – Input tensor
Returns:: Tensor with natural logarithm applied element-wise

Selection Operations¶

std::array<Tensor, 2> mllm::nn::functional::topk(const Tensor &x, int32_t k, int32_t dim = -1, bool largest = true, bool sorted = true)¶

Find the top-k values and their indices along a given dimension.

Parameters:

x – Input tensor
k – Number of top elements to retrieve
dim – Dimension along which to find top-k elements (default: -1, last dimension)
largest – Whether to return largest (true) or smallest (false) elements (default: true)
sorted – Whether to return elements in sorted order (default: true)

Returns:

Array containing values tensor and indices tensor

Element-wise Operations¶

Tensor mllm::nn::functional::clip(const Tensor &x, float min_val, float max_val)¶

Clip (limit) the values in a tensor to a specified range.

Parameters:

x – Input tensor
min_val – Minimum value
max_val – Maximum value

Returns:

Clipped tensor

Reduction Operations¶

Tensor mllm::nn::functional::min(const Tensor &x, int32_t dim = std::numeric_limits<int32_t>::max(), bool keep_dim = false)¶

Compute the minimum value of elements in the tensor.

Parameters:

x – Input tensor
dim – Dimension along which to compute minimum. If max int32_t, compute over all dimensions (default: max int32_t)
keep_dim – Whether to keep the reduced dimension (default: false)

Returns:

Tensor with minimum values

Tensor mllm::nn::functional::max(const Tensor &x, int32_t dim = std::numeric_limits<int32_t>::max(), bool keep_dim = false)¶

Compute the maximum value of elements in the tensor.

Parameters:

x – Input tensor
dim – Dimension along which to compute maximum. If max int32_t, compute over all dimensions (default: max int32_t)
keep_dim – Whether to keep the reduced dimension (default: false)

Returns:

Tensor with maximum values

Tensor mllm::nn::functional::sum(const Tensor &x, int32_t dim = std::numeric_limits<int32_t>::max(), bool keep_dim = false)¶

Compute the sum of elements in the tensor.

Parameters:

x – Input tensor
dim – Dimension along which to compute sum. If max int32_t, compute over all dimensions (default: max int32_t)
keep_dim – Whether to keep the reduced dimension (default: false)

Returns:

Tensor with sum values

Tensor mllm::nn::functional::mean(const Tensor &x, int32_t dim = std::numeric_limits<int32_t>::max(), bool keep_dim = false)¶

Compute the mean of elements in the tensor.

Parameters:

x – Input tensor
dim – Dimension along which to compute mean. If max int32_t, compute over all dimensions (default: max int32_t)
keep_dim – Whether to keep the reduced dimension (default: false)

Returns:

Tensor with mean values