oneflow

The oneflow package contains data structures for multi-dimensional tensors and defines mathematical operations over these tensors. Additionally, it provides many utilities for efficient serializing of Tensors and arbitrary types, and other useful utilities.

It has a CUDA counterpart, that enables you to run your tensor computations on an NVIDIA GPU with compute capability >= 3.0

Tensor

BoolTensor

Creates a Tensor with the dtype of oneflow.bool and the device on cpu, it has the same parameters as oneflow.Tensor()

ByteTensor

Creates a Tensor with the dtype of oneflow.uint8 and the device on cpu, it has the same parameters as oneflow.Tensor()

CharTensor

Creates a Tensor with the dtype of oneflow.int8 and the device on cpu, it has the same parameters as oneflow.Tensor()

DoubleTensor

Creates a Tensor with the dtype of oneflow.float64 and the device on cpu, it has the same parameters as oneflow.Tensor()

FloatTensor

Creates a Tensor with the dtype of oneflow.float32 and the device on cpu, it has the same parameters as oneflow.Tensor()

HalfTensor

Creates a Tensor with the dtype of oneflow.float16 and the device on cpu, it has the same parameters as oneflow.Tensor()

IntTensor

Creates a Tensor with the dtype of oneflow.int32 and the device on cpu, it has the same parameters as oneflow.Tensor()

LongTensor

Creates a Tensor with the dtype of oneflow.int64 and the device on cpu, it has the same parameters as oneflow.Tensor()

is_tensor

Note that this function is simply doing isinstance(obj, Tensor).

is_floating_point

Returns True if the data type of input is a floating point data type i.e., one of oneflow.float64 , oneflow.float32 , oneflow.float16, and oneflow.bfloat16.

is_nonzero

Returns True if the input is a single element tensor which is not equal to zero after type conversions.

numel

Returns the total number of elements in the input tensor.

set_printoptions

Set options for printing.

get_default_dtype

Returns the default floating point dtype.

set_default_dtype

Sets the default floating point type for those source operators which create Tensor.

set_default_tensor_type

Sets the default floating point type for those source operators which create Tensor.

Creation Ops

Note

Random sampling creation ops are listed under Random sampling and include: oneflow.rand() oneflow.randn() oneflow.randint() oneflow.randperm()

tensor

Constructs a tensor with data, return a global tensor if placement and sbp are in kwargs,

as_tensor

Converts data into a tensor, sharing data and preserving autograd history if possible.

as_strided

Create a view of an existing oneflow.Tensor input with specified size, stride and storage_offset.

from_numpy

Creates a Tensor from a numpy.ndarray.

zeros

Returns a tensor filled with the scalar value 0, with the shape defined by the variable argument size.

zeros_like

The interface is consistent with PyTorch.

ones

Returns a tensor filled with the scalar value 1, with the shape defined by the variable argument size.

ones_like

The interface is consistent with PyTorch.

randn_like

Returns a tensor with the same size as input that is filled with random numbers from a normal distribution with mean 0 and variance 1.

randint_like

Returns a tensor filled with random integers generated uniformly between low (inclusive) and high (exclusive).

masked_fill

Fills elements of self tensor with value where mask is True.

new_ones

The interface is consistent with PyTorch.

arange

Returns a 1-D tensor of size \(\left\lfloor \frac{\text{end} - \text{start}}{\text{step}} \right\rfloor + 1\) with values from start to end with step step.

linspace

Creates a one-dimensional tensor of size steps whose values are evenly spaced from start to end, inclusive.

eye

This operator creates a 2-D Tensor with ones on the diagonal and zeros elsewhere.

empty

The interface is consistent with PyTorch.

empty_like

The interface is consistent with PyTorch.

full

Creates a tensor of size size filled with fill_value.

full_like

Returns a tensor with the same size as input filled with fill_value.

tensor_scatter_nd_update

This operation creates a new tensor by applying sparse updates to the input tensor.

logspace

This function is equivalent to PyTorch’s logspace function.

Indexing, Slicing, Joining, Mutating Ops

argwhere

This operator finds the indices of input Tensor input elements that are non-zero.

atleast_1d

Returns a 1-dimensional view of each input tensor with zero dimensions.

atleast_2d

Returns a 2-dimensional view of each input tensor with zero dimensions.

atleast_3d

Returns a 3-dimensional view of each input tensor with zero dimensions.

cat

Concatenate two or more Tensor s at specified dim.

column_stack

Creates a new tensor by horizontally stacking the tensors in tensors.

concat

cat(tensors, dim=0) -> Tensor

chunk

Splits a tensor into a specific number of chunks.

dstack

Stack tensors in tensors depthwish (along third axis).

expand

This operator expand the input tensor to a larger size.

gather

Gathers values along an axis specified by dim.

gather_nd

This operator is a high-dimensional extension of gather, index is a K-dimensional tensor, which is regarded as a index of input Tensor input.

batch_gather

Gather the element in batch dims.

hsplit

The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.hsplit.html.

hstack

Stack tensors in tensors horizontally (column wise).

vsplit

Splits input, a tensor with two or more dimensions, into multiple tensors vertically according to indices_or_sections.

vstack

Stack tensors in tensors vertically (row wise).

index_select

Select values along an axis specified by dim.

index_add

See oneflow.Tensor.index_add_() for function description.

masked_select

Returns a new 1-D tensor which indexes the input tensor according to the boolean mask mask which is a BoolTensor(In oneFlow BoolTensor is replaced by Int8Tensor).

movedim

Moves the dimension(s) of input at the position(s) in source to the position(s) in destination.

narrow

Returns a new tensor that is a narrowed version of input tensor.

nonzero

permute

Returns a view of the original tensor with its dimensions permuted.

repeat

This operator repeat the input tensor to a larger size along the specified dimensions.

reshape

This operator reshapes a Tensor.

row_stack

Alias of oneflow.vstack().

select

Slices the self tensor along the selected dimension at the given index.

scatter

This operator writes the elements specified by index along with the axis dim from the src into the input.

scatter_add

This operator scatter the src with addition operation according to index along dim into the input.

scatter_nd

This operator inserts the elements in update according to the index and create a new Tensor.

slice

Extracts a slice from a tensor.

slice_update

Update a slice of tensor x.

split

Splits the tensor into chunks.

squeeze

This operator removes the specified dimention which size is 1 of the input Tensor.

stack

Concatenates a sequence of tensors along a new dimension.

swapaxes

This function is equivalent to NumPy’s swapaxes function.

swapdims

This function is equivalent to torch’s swapdims function.

t

oneflow.t(input) → Tensor.

tile

Constructs a tensor by repeating the elements of input.

transpose

Returns a tensor that is a transposed version of input.

unbind

Removes a tensor dimension.

unsqueeze

Returns a new tensor with a dimension of size one inserted at the specified position.

where

Return a tensor of elements selected from either x or y, depending on condition.

tensor_split

Splits a tensor into multiple sub-tensors, all of which are views of input, along dimension dim according to the indices or number of sections specified by indices_or_sections .

Random sampling

seed

Sets the seed for generating random numbers to a non-deterministic random number.

manual_seed

Sets the seed for generating random numbers.

initial_seed

Returns the initial seed for generating random numbers as a Python long.

get_rng_state

Sets the random number generator state.

set_rng_state

Returns the random number generator state as a oneflow.ByteTensor.

bernoulli

This operator returns a Tensor with binaray random numbers (0 / 1) from a Bernoulli distribution.

normal

Returns a tensor of random numbers drawn from separate normal distributions whose mean and standard deviation are given.

rand

Returns a tensor filled with random numbers from a uniform distribution on the interval [0, 1)

randint

Returns a tensor filled with random integers generated uniformly between low (inclusive) and high (exclusive).

randn

Returns a tensor filled with random numbers from a normal distribution with mean 0 and variance 1 (also called the standard normal distribution).

randperm

Returns a random permutation of integers from 0 to n - 1.

multinomial

Returns a tensor where each row contains num_samples indices sampled from the multinomial probability distribution located in the corresponding row of tensor input.

In-place random sampling

There are a few more in-place random sampling functions defined on Tensors as well. Click through to refer to their documentation: - oneflow.Tensor.normal_() - in-place version of oneflow.normal() - oneflow.Tensor.uniform_() - numbers sampled from the continuous uniform distribution

Serialization

save

Save an object to a directory.

load

Loads an object saved with oneflow.save() from a directory.

Parallelism

set_num_threads

Sets the number of threads used for intraop parallelism on CPU.

Locally disabling gradient computation

The context managers oneflow.no_grad(), oneflow.enable_grad(), and oneflow.set_grad_enabled() are helpful for locally disabling and enabling gradient computation. These context managers are thread local, so they won’t work if you send work to another thread using the threading module, etc.

Examples:

>>> import oneflow
>>> x = oneflow.zeros(1, requires_grad=True)
>>> with oneflow.no_grad():
...     y = x * 2
>>> y.requires_grad
False

>>> with oneflow.set_grad_enabled(False):
...     y = x * 2
>>> y.requires_grad
False

>>> with oneflow.set_grad_enabled(True):
...     y = x * 2
>>> y.requires_grad
True

no_grad

Context-manager that disabled gradient calculation.

set_grad_enabled

Context-manager that enabled gradient calculation.

enable_grad

Context-manager that enabled gradient calculation.

is_grad_enabled

Returns True if grad mode is currently enabled.

inference_mode

Context-manager that enables or disables inference mode

Math operations

Pointwise Ops

abs

Return the absolute value of each element in input tensor:math:y = |x| element-wise.

acos

Returns a new tensor with the inverse cosine of the elements of input.

acosh

Returns a new tensor with the inverse hyperbolic cosine of the elements of input.

arccos

Returns a new tensor with the inverse cosine of the elements of input.

arccosh

Returns a new tensor with the inverse hyperbolic cosine of the elements of input.

add

Adds other, scaled by alpha, to input.

addcdiv

This function is equivalent to PyTorch’s addcdiv function.

addcmul

Performs the element-wise multiplication of tensor1 by tensor2, multiply the result by the scalar value and add it to input.

asin

Returns a new tensor with the arcsine of the elements of input.

asinh

Returns a new tensor with the inverse hyperbolic sine of the elements of input.

arcsin

Returns a new tensor with the arcsine of the elements of input.

arcsinh

Returns a new tensor with the inverse hyperbolic sine of the elements of input.

atan

Returns a new tensor with the arctangent of the elements of input.

atanh

Returns a new tensor with the inverse hyperbolic tangent of the elements of input.

arctan

Returns a new tensor with the arctangent of the elements of input.

arctanh

Returns a new tensor with the inverse hyperbolic tangent of the elements of input.

atan2

Element-wise arctangent of input{i}/other{i} with consideration of the quadrant.

ceil

Returns a new tensor with the ceil of the elements of input, the smallest integer greater than or equal to each element.

ceil_

In-place version of oneflow.ceil()

clamp

Clamp all elements in input into the range [ min, max ] and return a resulting tensor:

clamp_min

Clamp all elements in input which are less than min to min and return a resulting tensor:

clamp_max

Clamp all elements in input which are greater than max to max and return a resulting tensor:

clip

Alias for oneflow.clamp().

cos

Returns a new tensor with the cosine of the elements of input.

cosh

Returns a new tensor with the hyperbolic cosine of the elements of input.

div

Computes the division of input by other for each element, scalar and broadcast promotation are supported.

erf

Computes the error function of each element.

erfc

Computes the complementary error function of each element of input.

erfinv

Computes the inverse error function of input.

exp

This operator computes the exponential of Tensor.

expm1

Returns a new tensor with the exponential of the elements minus 1 of input.

floor

Returns a new tensor with the arcsine of the elements of input.

floor_

In-place version of oneflow.floor()

frac

frac(input) → Tensor

frac_

In-place version of oneflow.frac().

fmod

Computes the element-wise remainder of division.

gelu

Applies the Gaussian Error Linear Units function:

quick_gelu

Applies GELU approximation that is fast but somewhat inaccurate.

square_relu

Applies the relu^2 activation introduced in https://arxiv.org/abs/2109.08668v2

log

Returns a new tensor with the natural logarithm of the elements of input.

log1p

Returns a new tensor with the natural logarithm of (1 + input).

log2

Returns a new tensor with the natural logarithm to the base 2 of the elements of input.

log10

Returns a new tensor with the natural logarithm to the base 10 of the elements of input.

logical_and

Computes the element-wise logical AND of the given input tensors.

logical_not

Computes the element-wise logical NOT of the given input tensors.

logical_or

Computes the element-wise logical OR of the given input tensors.

logical_xor

Computes the element-wise logical XOR of the given input tensors.

bitwise_and

Computes the bitwise AND of input and other.

bitwise_or

Computes the bitwise OR of input and other.

bitwise_xor

Computes the bitwise XOR of input and other.

bitwise_not

Computes the bitwise NOT of input.

mish

Applies the element-wise function:

mul

Computes the multiplication of input by other for each element, scalar and broadcast promotation are supported.

neg

This operator computes the negative value of Tensor.

negative

This operator computes the negative value of Tensor.

pow

Takes the power of each element in input with exponent and returns a tensor with the result.

reciprocal

Computes the safe reciprocal of x.

round

This operator rounds the value of Blob to the nearest integer.

round_

In-place version of oneflow.round().

rsqrt

Returns a new tensor with the reciprocal of the square-root of each of the elements of input.

selu

Applies element-wise function

softmax

Softmax is defined as:

softplus

Applies the element-wise function:

softsign

The formula is:

silu

The formula is:

sigmoid

Applies the element-wise function \(\text{Sigmoid}(x) = \frac{1}{1 + \exp(-x)}\)

sign

Computes the sign of Tensor.

sin

Returns a new tensor with the sine of the elements of input.

sinh

Returns a new tensor with the hyperbolic sine of the elements of input.

sin_

In-place version of oneflow.sin()

sqrt

Returns a new tensor with the square-root of the elements of input.

square

Returns a new tensor with the square of the elements of input.

sub

Computes the subtraction of input by other for each element, scalar and broadcast promotation are supported.

tan

Returns the tan value of the elements of input.

tanh

The equation is:

trunc

The interface is consistent with PyTorch.

floor_divide

lerp

The documentation is referenced from: https://pytorch.org/docs/stable/generated/torch.lerp.html.

lerp_

In-place version of oneflow.lerp()

quantile

The documentation is referenced from: https://pytorch.org/docs/stable/generated/torch.quantile.html.

Reduction Ops

argmax

The op computes the index with the largest value of a Tensor at specified axis.

argmin

The op computes the index with the largest value of a Tensor at specified axis.

amax

Returns the maximum along a dimension.

amin

Returns the minimum value of each slice of the input tensor in the given dimension(s) dim.

any

For each row of input in the given dimension dim, returns True if any element in the row evaluate to True and False otherwise.

max

Computes the maximum value of all elements in the input tensor.

min

Computes the minimum value of all elements in the input tensor.

mean

Computes the mean of row of elements in a tensor in the given dimension.

median

Returns the median of the values in input.

mode

Returns a namedtuple (values, indices) where values is the mode value of each row of the input tensor in the given dimension dim, i.e. a value which appears most often in that row, and indices is the index location of each mode value found.

prod

Computes the product of row of elements in a tensor in the given dimension.

nansum

Returns the sum of each row of the input tensor in the given dimension dim, treating Not a Numbers (NaNs) as zero.

std

Returns the standard-deviation of each row of the input tensor in the dimension dim.

sum

Computes the sum of row of elements in a tensor in the given dimension.

logsumexp

Returns the log of summed exponentials of each row of the input tensor in the given dimension dim.

var

Returns the variance of each row of the input tensor in the given dimension dim.

norm

Returns the matrix norm or vector norm of a given tensor.

all

For each row of input in the given dimension dim, returns True if all element in the row evaluate to True and False otherwise.

Comparison Ops

argsort

This operator sorts the input Tensor at specified dim and returns the indices of the sorted Tensor.

eq

Computes element-wise equality.

equal

True if two tensors have the same size and elements, False otherwise.

gt

Returns the truth value of \(input > other\) element-wise.

isinf

This function is equivalent to PyTorch’s isinf function.

isnan

This function is equivalent to PyTorch’s isnan function.

le

Returns the truth value of \(input <= other\) element-wise.

lt

Returns the truth value of \(input < other\) element-wise.

ne

Computes element-wise not equality.

sort

Sorts the elements of the input tensor along a given dimension in ascending order by value.

topk

Finds the values and indices of the k largest entries at specified axis.

ge

Returns the truth value of \(input >= other\) element-wise.

greater

Returns the truth value of \(input > other\) element-wise.

greater_equal

Returns the truth value of \(input >= other\) element-wise.

maximum

Computes the element-wise maximum of x and y.

minimum

Computes the element-wise minimum of x and y.

not_equal

ne(input, other) -> Tensor

isclose

The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.isclose.html

allclose

The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.allclose.html

Spectral Ops

hann_window

This function is equivalent to PyTorch’s hann_window function.

Other Ops

adaptive_avg_pool1d

Applies a 1D adaptive average pooling over an input signal composed of several input planes.

adaptive_avg_pool2d

Applies a 2D adaptive average pooling over an input signal composed of several input planes.

adaptive_avg_pool3d

Applies a 3D adaptive average pooling over an input signal composed of several input planes.

broadcast_like

This operator broadcast tensor x to like_tensor according to the broadcast_axes.

cast

The operation takes input tensor x and casts it to the output with dtype

cumprod

This operator computes the cumulative product of input elements in the given dimension.

cumsum

This operator computes the cumulative sum of input elements in the given dimension.

diag

If input is a vector (1-D tensor), then returns a 2-D square tensor with the elements of input as the diagonal.

diagonal

Returns a partial view of input with the its diagonal elements with respect to dim1 and dim2 appended as a dimension at the end of the shape.

einsum

Sums the product of the elements of the input operands along dimensions specified using a notation based on the Einstein summation convention.

flatten

Flattens a contiguous range of dims into a tensor.

flip

Reverse the order of a n-D tensor along given axis in dims.

in_top_k

Says whether the targets are in the top K predictions.

meshgrid

Take \(N\) tensors, each of which can be either scalar or 1-dimensional vector, and create \(N\) N-dimensional grids, where the \(i\) th grid is defined by expanding the \(i\) th input over dimensions defined by other inputs.

nms

Performs non-maximum suppression (NMS) on the boxes according to their intersection-over-union (IoU).

roc_auc_score

roll

Roll the tensor along the given dimension(s).

searchsorted

Find the indices from the innermost dimension of sorted_sequence such that, if the corresponding values in values were inserted before the indices, the order of the corresponding innermost dimension within sorted_sequence would be preserved.

tensordot

Compute tensor dot along given dimensions.

tril

Returns the lower triangular part of a matrix (2-D tensor) or batch of matrices input along the specified diagonal, the other elements of the result tensor out are set to 0.

repeat_interleave

Repeat elements of a tensor.

triu

Returns the upper triangular part of a matrix (2-D tensor) or batch of matrices input, the other elements of the result tensor out are set to 0.

cross

Returns the cross product of vectors in dimension dim of input and other.

bincount

oneflow.bincount(input, weights=None, minlength=0) → Tensor

broadcast_shapes

The interface is consistent with PyTorch.

broadcast_tensors

The interface is consistent with PyTorch.

broadcast_to

The interface is consistent with PyTorch.

unique

Returns the unique elements of the input tensor.

BLAS and LAPACK Operations

addmm

Performs a matrix multiplication of the matrices mat1 and mat2.

bmm

Performs a batch matrix-matrix product of matrices stored in input and mat2.

baddbmm

The documentation is referenced from: https://pytorch.org/docs/1.10/generated/torch.baddbmm.html.

dot

This operator computes the dot product of tensor input and other.

matmul

This operator applies matrix multiplication to two Tensor.

mm

Performs a matrix multiplication of the matrices input and mat2.

mv

Performs a matrix-vector product of the matrix input and the vector vec.