Module ndarray

Module ndarray 

Source
Expand description

§N-dimensional Array for Linear Algebra & Tensor Computations

An NdArray is a fixed-size multidimensional array container defined by its shape and datatype. 1D (vectors) and 2D (matrices) arrays are often of special interest and can be used in various linear algebra computations including dot products, matrix products, batch matrix multiplications, einsums, and more.

NdArrays can be iterated over (along configurable dimensions), reshaped, sliced, and indexed, reduced, and more.

This struct is heavily modeled after NumPy’s ndarray and supports many of the same methods.

Example:

use redstone_ml::*;

let matrix_a = NdArray::new([[1, 3, 2], [-1, 0, -1]]); // shape [2, 3]
let matrix_b = NdArray::randint([3, 7], -5, 3);

let matrix_view = matrix_b.slice_along(Axis(1), 0..2); // shape [3, 2]
let matrix_c = matrix_a.matmul(matrix_view);

let result = matrix_c.sum();

§NdArray Views & Lifetimes

There are 2 ways we can create NdArray views: by borrowing or by consuming:

let data = NdArray::<f64>::rand([9]);
let matrix = (&data).reshape([3, 3]); // by borrowing (data remains alive after)

let data = NdArray::<f64>::rand([9]);
let matrix = data.reshape([3, 3]); // by consuming data

The consuming syntax allows us to chain operations without worrying about lifetimes

// a reshaped and transposed random matrix
let matrix = NdArray::<f64>::rand([9]).reshape([3, 3]).T();

Operations like reshape, view, diagonal, squeeze, unsqueeze, T, transpose, and ravel do not create new NdArrays by duplicating memory (which would be slow). They always return NdArray views which share memory with the source NdArray. NdArray::clone() or NdArray::flatten() can be used to duplicate the underlying NdArray.

This means that all NdArray views have a lifetime at-most as long as the source NdArray.

§Linear Algebra, Broadcasting, and Reductions

We currently support the core linear algebra operations including dot products, matrix-vector and matrix-matrix multiplications, batched matrix multiplications, and trace.

vector1.dot(vector2);

matrix.trace(); // also trace_along/offset_trace
matrix.diagonal(); // also diagonal_along/offset_diagonal
matrix.matmul(vector);
matrix1.matmul(matrix2);

batch_matrices1.bmm(batch_matrices2);

We can also perform various reductions including sum, product, min, max, min_magnitude, and max_magnitude. Each of these is accelerated with various libraries including vDSP, Arm64 NEON SIMD, and BLAS.

let sum = ndarray.sum();
let sum_along = ndarray.sum_along([0, -1]); // sum along first and last axes

NdArrays can be used in arithmetic operations using the usual binary operators including addition (+), subtraction (-), multiplication (*), division (/), remainder (%), and bitwise operations (&, |, <<, >>).

let result = &arr1 + &arr2; // non-consuming
let result = &arr1 + arr2;  // consumes RHS
let result = arr1 + arr2;   // consumes both

NdArrays are automatically broadcast using the exact same rules as NumPy to perform efficient computations with different-dimensional (yet compatible) data.

§Slicing, Indexing, and Iterating

Slicing and indexing an NdArray always return a view. This is how we can access various elements of vectors, columns/rows of matrices, and more.

let arr = NdArray::<f32>::rand([2, 4, 3, 5]); // 4D NdArray
let slice1 = arr.slice(s![.., 0, ..=2]);      // use s! to specify a slice
let slice2 = arr.slice_along(Axis(-2), 0);    // 0th element along second-to-last axis
let el = arr[[0, 3, 2, 4]];

One can also iterate over an NdArray in various ways:

for subarray in arr.iter() { /* 4x3x5 subarrays */ }
for subarray in arr.iter_along(Axis(2)) { /* 2x4x5 subarrays */ }
for el in arr.flatiter() { /* element-wise iteration */ }

§Other Constructors

let ndarray = NdArray::arange(0i32, 5); // [0, 1, 2, 3, 4]
let ndarray = NdArray::linspace(0f32, 1.0, 5); // [0.0, 0.25, 0.5, 0.75, 1.0]
let ndarray = NdArray::full(5.0, [5, 4, 2]);
let falses = NdArray::<bool>::zeros([5, 4, 2]);

A scalar NdArray is dimensionless and contains a single value. It is often the return value for reduction methods like sum, product, min, and max.

let ten = NdArray::scalar(10u8);

In many cases, one desires randomized multidimensional arrays with a specified shape.

let rand = NdArray::<f32>::randn([3, 4]);
let rand = NdArray::<f32>::rand([3, 4]);
let rand = NdArray::randint([3, 4], -5, 3);

Re-exports§

pub use iterator::*;

Modules§

astype
binary_ops
broadcast
clone
constructors
equals
fill
index_impl
iterator
methods
reduce
reshape
slice

Structs§

NdArray