Array

Array#

Core Array class with improved organization.

class nabla.core.array.Array(shape, dtype=<MagicMock name='mock.float32' id='140511747382336'>, device=max.driver.CPU, materialize=False, name='', batch_dims=())[source]#

Bases: object

Core tensor-like array class with automatic differentiation support.

__init__(shape, dtype=<MagicMock name='mock.float32' id='140511747382336'>, device=max.driver.CPU, materialize=False, name='', batch_dims=())[source]#

shape: tuple[int, ...]#

batch_dims: tuple[int, ...]#

dtype: <MagicMock id='140511747387904'>#

device: max.driver.Device#

name: str#

args: list[Array]#

visited: bool#

tensor_value: <MagicMock id='140511747378352'> | <MagicMock id='140511747356960'> | <MagicMock id='140511747358256'> | None#

maxpr: Array], None] | None#

vjp_rule: Callable[[list[Array], Array, Array], list[Array]] | None#

jvp_rule: Callable[[list[Array], list[Array], Array], Array] | None#

traced: bool#

tangent: Array | None#

cotangent: Array | None#

stage_realization: bool#

kernel_impl_path: Path | None#

custom_kernel_path: Path | None#

property impl: max.driver.Tensor | None#: Get the max.Tensor representation of this Array. If the underlying _impl field is a Numpy array, convert it to a Tensor.

impl_(value)[source]#

Set the implementation of this Array to a Numpy array or Tensor.

property size: int#: Return the total number of elements in the array.

classmethod from_impl(impl, name='')[source]#

Create Array from existing Tensor implementation.

copy_from(other)[source]#

Copy data from another Array.

add_arguments(*arg_nodes)[source]#

Add an arguments to this Array’s computation graph if traced.

realize()[source]#

Force computation of this Array.

to_numpy()[source]#

Get NumPy representation.

classmethod from_numpy(np_array)[source]#

Create a new Array from a NumPy array.

get_arguments()[source]#

Get list of argument Arrays.

set_maxpr(fn)[source]#

Set the MAX PR function for this operation.

to(device)[source]#

Move Array to specified device.

astype(dtype)[source]#

Convert array to a different data type.

Parameters:: dtype (<MagicMock id='140511747387904'>) – Target data type
Returns:: New Array with the specified data type
Return type:: Array

sum(axes=None, keep_dims=False)[source]#

Sum array elements over given axes.

Parameters:

axes – Axis or axes along which to sum. Can be int, list of ints, or None (sum all)
keep_dims – If True, reduced axes are left as dimensions with size 1

Returns:

Array with the sum along the specified axes

Return type:

Array

Examples:

arr.sum()           # Sum all elements
arr.sum(axis=0)     # Sum along first axis
arr.sum(axis=[0,1]) # Sum along first two axes

reshape(shape)[source]#

Change the shape of an array without changing its data.

Parameters:: shape (tuple[int, ...]) – New shape for the array
Returns:: Array with the new shape
Return type:: Array

Examples:

arr.reshape((2, 3))     # Reshape to 2x3
arr.reshape((-1,))      # Flatten to 1D (note: -1 not yet supported)

permute(axes)[source]#

Permute the dimensions of the array.

Parameters:: axes (tuple[int, ...]) – List of integers specifying the new order of dimensions
Returns:: Array with dimensions permuted according to the specified axes
Return type:: Array

Examples:

arr.permute([1, 0]) # If arr.shape is (2, 3), this will return an array with shape (3, 2)

transpose(axes)[source]#

Permute the dimensions of the array.

Parameters:: axes (tuple[int, ...]) – List of integers specifying the new order of dimensions
Returns:: Array with dimensions permuted according to the specified axes
Return type:: Array

Examples:

arr.permute([1, 0]) # If arr.shape is (2, 3), this will return an array with shape (3, 2)

at(key, value)[source]#: Update array at specified indices/slices, returning new array.

set(key, value)[source]#

Set values at specified indices/slices, returning a new array.

This is a functional operation that returns a new Array with the specified values updated, leaving the original Array unchanged.

Parameters:

key – Index specification (int, slice, tuple of indices/slices, ellipsis)
value – Value(s) to set at the specified location

Returns:

New Array with updated values

Return type:

Array

Examples

new_arr = arr.set(1, 99.0) # Set single element new_arr = arr.set((1, 2), 99.0) # Set element at (1,2) new_arr = arr.set(slice(1, 3), 99.0) # Set slice new_arr = arr.set(…, 99.0) # Set with ellipsis

Array

Contents

Array#