megengine.functional

import megengine.functional as F

Note

顾名思义，megengine.functional 模块中包含着所有与 Tensor 有关的计算接口：

• 与神经网络（Neural Network）相关的算子统一封装在 megengine.functional.nn 中；

• 分布式算子统一封装在 megengine.functional.distributed 中，方便调用；

• 其它的常见算子均可在 megengine.functional 中直接调用；

See also

用户指南中对于 Use Functional operations and calculations 有另外一套分类逻辑，可作为参考。

General tensor operations

Note

该部分的 API 设计接纳了 Python 数据 API 标准联盟 中的倡导, 尽可能地向 NumPy API 靠拢。

Creation Functions

See also

How to create a Tensor

arange	Returns evenly spaced values within the half-open interval [start, stop) as a one-dimensional tensor.
linspace	Returns evenly spaced numbers over a specified interval.
eye	Returns a two-dimensional tensor with ones on the diagonal and zeros elsewhere.
diag	Extract a diagonal or construct a diagonal tensor.
zeros	Returns a new tensor having a specified shape and filled with zeros.
zeros_like	Returns a tensor filled with zeros with the same shape and data type as input tensor.
ones	Returns a new tensor having a specified shape and filled with ones.
ones_like	Returns a tensor filled with ones with the same shape and data type as input tensor.
full	Returns a new tensor having a specified shape and filled with given value.
full_like	Returns a tensor filled with given value with the same shape as input tensor.
complex	Constructs a complex tensor with its real part equal to real and its imaginary part equal to imag.
polar	Constructs a complex tensor whose elements are Cartesian coordinates corresponding to the polar coordinates with absolute value abs and angle angle.

Manipulation Functions

See also

How to operate Tensor

copy	Copies tensor to another device.
reshape	Reshapes a tensor without changing its data.
flatten	Reshapes the tensor by flattening the sub-tensor from dimension start_axis to dimension end_axis.
transpose	Swaps shapes and strides according to given pattern.
swapaxes	Interchange two axes of a tensor.
broadcast_to	Broadcasts a tensor to given shape.
expand_dims	Adds dimension before given axis.
squeeze	Removes dimension of shape 1.
concat	Concat some tensors
stack	Concats a sequence of tensors along a new axis.
split	Splits the input tensor into several smaller tensors.
tile	Construct an array by repeating inp the number of times given by reps.
repeat	Repeat elements of an array.
roll	Roll the tensor along the given axis(or axes).
real	Returns a new tensor containing real values of the complex tensor.
imag	Returns a new tensor containing imaginary values of the complex tensor.
cross	Return the cross product of two (arrays of) vectors.

Arithmetic operations

add	Element-wise addition.
sub	Element-wise subtraction.
mul	Element-wise multiplication.
div	Element-wise division.
floor_div	Element-wise floor division.
neg	Element-wise negation.
abs	Element-wise \(\operatorname{abs}(x)\) function.
sign	Element-wise sign.
pow	Element-wise power.
mod	Element-wise \(\operatorname{mod}(x, y)\) function.
sqrt	Element-wise \(\operatorname{sqrt}(x)\) function.
square	Element-wise \(x^2\) function.
maximum	Element-wise maximum of tensor elements.
minimum	Element-wise minimum of tensor elements.
round	Element-wise \(\operatorname{round}(x)\) function.
ceil	Element-wise \(\lceil x \rceil\) function.
floor	Element-wise \(\lfloor x \rfloor\) function.
clip	Element-wise clipping function.
exp	Element-wise \(e^x\) function.
expm1	Element-wise \(e^x - 1\) function.
log	Element-wise \(\operatorname{log}(x)\) function.
log1p	Element-wise \(\log(1 + x)\) function.
logaddexp	Element-wise \(\log(e^x + e^y)\) function.

Trigonometric functions

sin	Element-wise \(\sin(x)\) function.
cos	Element-wise \(\cos(x)\) function.
tan	Element-wise \(\tan(x)\) function.
asin	Element-wise \(\arcsin(x)\) function.
acos	Element-wise \(\arccos(x)\) function.
atan	Element-wise \(\arctan(x)\) function.
atan2	Element-wise \(\arctan(\frac{y}{x})\) function.

Hyperbolic functions

sinh	Element-wise \(\sinh(x)\) function.
cosh	Element-wise \(\cosh(x)\) function.
tanh	Element-wise \(\tanh(x)\) function.
acosh	Element-wise \(\cosh^{-1}(x)\) function.
asinh	Element-wise \(\sinh^{-1}(x)\) function.
atanh	Element-wise \(\tanh^{-1}(x)\) function.

Bit operations

left_shift	Element-wise left shift.
right_shift	Element-wise right shift.

Logic functions

isnan	Element-wise NaN check.
isinf	Element-wise infinity check.
logical_and	Element-wise logical AND.
logical_not	Element-wise logical NOT.
logical_or	Element-wise logical OR.
logical_xor	Element-wise logical XOR.
greater	Element-wise greater-than comparison.
greater_equal	Element-wise greater-than-or-equal-to comparison.
less	Element-wise less-than comparison.
less_equal	Element-wise less-than-or-equal-to comparison.
equal	Element-wise equality comparison.
not_equal	Element-wise inequality comparison.

Statistical Functions

sum	Calculates the sum of tensor elements over a given axis (or axes).
cumsum	Calculates the cumulative sum of tensor elements over a given axis.
prod	Calculates the product of tensor elements over a given axis (or axes).
mean	Calculates the mean of tensor elements over a given axis (or axes).
min	Calculates the minimum of tensor elements over a given axis (or axes).
max	Calculates the maximum of tensor elements over a given axis (or axes).
var	Calculates the variance of tensor elements over a given axis (or axes).
std	Calculates the standard deviation of tensor elements over a given axis (or axes).
norm	Calculates the norm of tensor elements over a given axis.
normalize	Performs \(L_p\) normalization of input tensor along given axis.

See also

想要返回 min, max 的索引而不是元素值，请参考 Searching Functions

Linear Algebra Functions

dot	Computes dot-product of two vectors inp1 and inp2.
matinv	Computes the inverse of a batch of matrices; input must has shape [..., n, n].
matmul	Performs a matrix multiplication of the matrices inp1 and inp2.
svd	Computes the singular value decomposition of a matrix (or a stack of matrices) inp.

Indexing Functions

gather	Gathers data from input tensor on axis using index.
scatter	Writes all values from the tensor source into input tensor at the indices specified in the index tensor.
cond_take	Takes elements from data if specific condition is satisfied on mask.
where	Selects elements either from Tensor x or Tensor y, according to mask.
non_zero	When as_tuple is False (default): Returns a tensor including the indices of all non-zero elements of Tensor condition.

Searching Functions

argmin	Returns the indices of the minimum values along given axis.
argmax	Returns the indices of the maximum values along given axis.

Sorting Functions

sort	Returns sorted tensor and the indices would sort the input tensor.
argsort	Returns the indices that would sort the input tensor.
topk	Selects the Top-K (by default) smallest elements of 2d matrix by row.

Metric functions

topk_accuracy

Calculates the classification accuracy given predicted logits and ground-truth labels.

Neural network operations

Convolution functions

conv1d	1D convolution operation.
conv2d	2D convolution operation.
conv3d	3D convolution operation.
local_conv2d	Applies a spatial convolution with untied kernels over an groupped channeled input 4D tensor.
conv_transpose2d	2D transposed convolution operation.
conv_transpose3d	3D transposed convolution operation.
deformable_conv2d	Deformable Convolution.
sliding_window	Extracts sliding local blocks from a batched input tensor.
sliding_window_transpose	Sum over the sliding windows on the corresponding input location.

Pooling functions

avg_pool2d	Applies 2D average pooling over an input tensor.
max_pool2d	Applies a 2D max pooling over an input tensor.
adaptive_avg_pool2d	Applies a 2D average adaptive pooling over an input.
adaptive_max_pool2d	Applies a 2D max adaptive pooling over an input.
deformable_psroi_pooling	Deformable PSROI(Position Sensitive Region of Interest) Pooling.

Non-linear activation functions

sigmoid	Element-wise 1 / ( 1 + exp( -x ) ).
hsigmoid	Element-wise relu6(x + 3) / 6.
hswish	Element-wise x * relu6(x + 3) / 6.
relu	Applies the rectified linear unit function element-wise.
relu6	Element-wise min(max(x, 0), 6).
prelu	Element-wise max(x, 0) + y * min(x, 0).
leaky_relu	Element-wise LeakyReLU function
silu	Applies the element-wise Sigmoid Linear Unit function, i.e. x * sigmoid(x).
gelu	Applies the element-wise function:
softplus	Applies the element-wise function:
softmax	Applies a \(\text{softmax}(x)\) function.
logsoftmax	Applies the \(\log(\text{softmax}(x))\) function to an n-dimensional input tensor.
logsigmoid	Applies the element-wise function:
logsumexp	Calculates the logarithm of the inputs' exponential sum along the given axis.

Normalization functions

batch_norm	Applies batch normalization to the input.
sync_batch_norm	Applies synchronized batch normalization to the input.
layer_norm	Applies layer normalization to the input.
local_response_norm	Apply local response normalization to the input tensor.

Linear functions

linear

Applies a linear transformation to the input tensor.

multi-head attention functions

multi_head_attention

Allows the model to jointly attend to information from different representation subspaces.

Dropout functions

dropout

Returns a new tensor where each of the elements are randomly set to zero with probability P = drop_prob.

Sparse functions

one_hot	Performs one-hot encoding for the input tensor.
indexing_one_hot	One-hot indexing for some axes.
embedding	Applies lookup table for embedding.

Loss functions

l1_loss	Calculates the mean absolute error (MAE) between each element in the pred \(x\) and label \(y\).
square_loss	Calculates the mean squared error (squared L2 norm) between each element in the pred \(x\) and label \(y\).
hinge_loss	Caculates the hinge loss which is often used in SVM.
binary_cross_entropy	Computes the binary cross entropy loss (using logits by default).
cross_entropy	Computes the multi-class cross entropy loss (using logits by default).
ctc_loss	The Connectionist Temporal Classification loss.

Vision functions

cvt_color	Convert images from one format to another
pad	Pads the input tensor.
pixel_shuffle	Rearranges elements in a tensor of shape (..., C * r^2, H, W) to a tensor of shape (..., C, H * r, W * r), where r is an upscale factor, where ... is zero or more batch dimensions.
interpolate	Down/up samples the input tensor to either the given size or with the given scale_factor.
remap	Applies remap transformation to batched 2D images.
warp_affine	Batched affine transformation on 2D images.
warp_perspective	Applies perspective transformation to batched 2D images.
roi_pooling	Applies RoI (Region of Interest) pooling on input feature, as described in Faster RCNN.
roi_align	Applies RoI (Region of Interest) align on input feature, as described in Mask R-CNN.
nms	Performs non-maximum suppression (NMS) on the boxes according to their intersection-over-union(IoU).
correlation	Applies correlation to inputs.
nvof	Implements NVIDIA Optical Flow SDK.

Quantized functions

conv_bias_activation	Convolution bias with activation operation, only for inference.
batch_conv_bias_activation	Batch convolution bias with activation operation, only for inference.
conv_transpose2d	rtype Tensor

Distributed functions

all_gather	Gather tensors across the specified group and concat them at first dimension.
all_reduce_max	Reduce tensors with max operation on each value across the specified group.
all_reduce_min	Reduce tensors with min operation on each value across the specified group.
all_reduce_sum	Reduce tensors with sum operation on each value across the specified group.
all_to_all	Each process scatter input tensor to all processes and return gathered tensor.
broadcast	Broadcast tensor data from root process to others.
reduce_scatter_sum	Reduce tensors across the specified group by sum and split them at first dimension.
collective_comm	Helper function for applying collective communication functions.
remote_send	Send tensor to another process.
remote_recv	Receive a tensor from another process.
reduce_sum	Reduce tensor data across the specified group by sum.
gather	Gather tensors across the specified group.
scatter	Split tensor in root process at first dimension.

External functions

extern_opr_subgraph	Load a serialized extern opr subgraph and fake execute the operator.
tensorrt_runtime_opr
cambricon_runtime_opr	Load a serialized Cambricon model as a runtime operator in MegEngine.
atlas_runtime_opr	Load a serialized Atlas model as a runtime operator in MegEngine.
magicmind_runtime_opr	Load a serialized MagicMind model as a runtime operator in MegEngine.

previous

Collator

next

megengine.functional.arange