Module dfdx::tensor_ops

Expand description

Implementations of all ops for tensors including activations like relu(), binary operations like matmul(), and more.

Generic function and struct methods

All functionality is provided in two ways.

The generic standalone function that takes a generic parameter. e.g. mean().
The struct method for tensor structs. e.g. crate::prelude::Tensor1D::mean().

The struct methods are all just pass throughs to the generic function.

Axes/Dimensions for broadcasting/reductions/selecting

For the following sections, some traits/functions utilizing const isize to determine the axis to apply the transformation to.

Here are the valid axes for each tensor:

Tensor0D: -1
Tensor1D: 0, -1
Tensor2D: 0, 1, -1
Tensor3D: 0, 1, 2, -1
Tensor4D: 0, 1, 2, 3, -1

Note that -1 must be used for the last axis instead of the positive number. This is to prevent ambiguous operations.

Reductions

There are a number of functions that reduce a dimension (e.g. mean_axis()). These functions are all labeled with *_axis() at the end.

Reducing a dimension means removing that dimension from the tensor by reducing it to 1 number. For example calling sum_axis::<-1>() on a Tensor2D<2, 5> would result in a Tensor1D<2>. Calling sum_axis::<0>() on a Tensor2D<5> would result in a Tensor1D<5>.

See Reduce1 implementations for a complete list of reductions.

See relevant functions for more examples.

Broadcasts

Broadcasting tensors is provided through the following traits:

Broadcast1::broadcast1(), which broadcasts a single axis
Broadcast2::broadcast2(), which broadcasts 2 axes
Broadcast3::broadcast3(), which broadcasts 3 axes
Broadcast4::broadcast4(), which broadcasts 4 axes

See the implementations of each trait for a complete list of possible broadcasts.

To broadcast a tensor to be the same size as another tensor you can use like so:

let big: Tensor2D<2, 5> = TensorCreator::zeros();

// broadcast the 1nd axis
let a: Tensor2D<2, 5> = Tensor1D::<5>::zeros().broadcast1();
add(a, &big);

// broadcast the 2nd axis
let a: Tensor2D<2, 5> = Tensor1D::<2>::zeros().broadcast1();
add(a, &big);

Selects/Indexing

Selecting or indexing into a tensor is done via Select1::select(). This traits enables 2 behaviors for each axis of a given tensor:

Select exactly 1 element from that axis.
Select Z elements (can be different from the size of the axis) from that axis

For example here is selecting from the 0th axis of a 2d tensor:

let t = Tensor2D::new([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);

let a: Tensor1D<3> = t.clone().select(&0); // select the first row
assert_eq!(a.data(), &[1.0, 2.0, 3.0]);

let b: Tensor2D<5, 3> = t.select(&[0, 0, 1, 1, 1]); // select each row multiple times

This can be done per axis as well, which allows a number of combinations. Here is the same example but selecting from the last axis of a 2d tensor:

let t = Tensor2D::new([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]);

let a: Tensor1D<2> = t.clone().select(&[0, 2]); // select one element from the last axis
assert_eq!(a.data(), &[1.0, 6.0]);

let b: Tensor2D<2, 2> = t.select(&[[0, 2], [1, 1]]); // select multiple from the last axis
assert_eq!(b.data(), &[[1.0, 3.0], [5.0, 5.0]]);

Re-exports

pub use binary_map::*;

Modules

binary_map

Traits

Broadcast1

Broadcasts the Ith dimension. Increases number dimensions by 1. Results in T. Opposite of Reduce1.

Broadcast2

Broadcasts dimensions I1 and I2. Increases number dimensions by 2. Results in T.

Broadcast3

Broadcasts dimensions I1, I2, and I3. Increases number dimensions by 3. Results in T.

Broadcast4

Broadcasts dimensions I1, I2, I3, and I4. Increases number dimensions by 4. Results in T.

MatMulTrTyping

Enables concrete output types for generic matmul functions. Without this you’d have to specify type of output.

MatMulTyping

Enables concrete output types for generic matmul functions. Without this you’d have to specify type of output.

Reduce1

Reduce the Ith dimension of a Tensor. Enables functions like sum_axis() that reduce values along a single dimension.

Reshape

Requires Nightly Reshape Self into T.

Select1

Select values along a single axis I resulting in T. Equivalent to torch.select and torch.gather from pytorch.

Functions

abs

Absolute value (abs). |t|

add_scalar

t + val. val is used for all elements of t.

backward

Runs backprop algorithm with all operations contained in the tape that t has.

clamp

Clamp all elements between the provided min and max values.

conv2d

Requires Nightly Perform a 2d convolution.

conv2d_batched

Requires Nightly Perform a batched 2d convolution

cos

Cosine function.

div_scalar

t / val. val is used for all elements of t.

dropout

Does nothing if no tape is in t. Zeros elements with probability p and scales all elements by 1 / (1 - p). See Tape::OWNS_TAPE.

exp

Exponential function (exp). e^t

ln

Natural Logarithm (ln). log_e(t).

log_softmax

log(softmax(t)) in numerically stable way. Does t - logsumexp(t) under the hood.

logsumexp

Computes the LogSumExp function.

map

f(t). Applies a function f to every element of the Tensor. The derivative df must also be provided.

map_df_uses_fx

Same as map(), but calls df with the result of f(x). This can potentially remove an allocation.

matmul

Matrix multiplication. This also supports batched matrix multiplication, and broadcasted matrix multiplication.

matmul_transpose

Matrix multiplication with the transpose of rhs. Equivalent to matmul(lhs, transpose(rhs)). This supports the same variants as matmul (broadcasted, batched, etc).

max_axis

Reduces dimension I of the tensor by gathering the maximum value from that dimension.

mean

Average of all elements in the tensor. Returns a Tensor0D (i.e. one number).

mean_axis

Average the values along dimension I of T.

min_axis

Reduces dimension I of the tensor by gathering the minimum value from that dimension.

mul_scalar

t * val. val is used for all elements of t.

nans_to

Replaces any std::f32::NAN with value.

negate

Negates all elements.

normalize_axis

Normalizes t to have mean 0.0 and stddev 1.0 along the I dimension of T. epsilon is passed to std_axis(). Computes (t - t.mean(I)) / t.std(I, epsilon).

relu

Rectified Linear Unit (ReLU). max(0, t)

Computes the softmax function. Equivalent to exp(log_softmax(t)).

sqrt

√t or t^0.5

square

t^2