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 dataThe 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 axesNdArrays 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 bothNdArrays 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::*;