adam.core.array_api_math

adam.core.array_api_math#

Classes#

ArraySpec

ArrayAPILike

Generic Array-API-style wrapper used by NumPy/JAX/Torch backends.

ArrayAPIFactory

Generic factory. Give it (a) a Like class and (b) an xp namespace

ArrayAPISpatialMath

A drop-in SpatialMath that implements sin/cos/outer/concat/skew with the Array API.

Functions#

spec_from_reference(→ ArraySpec)

xp_getter(*xs)

Module Contents#

class adam.core.array_api_math.ArraySpec[source]#
xp: types.ModuleType[source]#
dtype: Any | None[source]#
device: Any | None[source]#
adam.core.array_api_math.spec_from_reference(ref: Any) → ArraySpec[source]#
adam.core.array_api_math.xp_getter(*xs: Any)[source]#
class adam.core.array_api_math.ArrayAPILike[source]#

Bases: adam.core.spatial_math.ArrayLike

Generic Array-API-style wrapper used by NumPy/JAX/Torch backends.

array: Any[source]#
__getitem__(idx) → ArrayAPILike[source]#
property shape[source]#
reshape(*args)[source]#
property T: ArrayAPILike[source]#

Transpose of the array

Type:

Returns

__matmul__(other)[source]#
__rmatmul__(other) → ArrayAPILike[source]#
__mul__(other) → ArrayAPILike[source]#
__rmul__(other) → ArrayAPILike[source]#
__truediv__(other) → ArrayAPILike[source]#
__add__(other) → ArrayAPILike[source]#
__radd__(other) → ArrayAPILike[source]#
__sub__(other) → ArrayAPILike[source]#
__rsub__(other) → ArrayAPILike[source]#
__neg__() → ArrayAPILike[source]#
property ndim[source]#
class adam.core.array_api_math.ArrayAPIFactory(like_cls, xp, *, dtype=None, device=None)[source]#

Bases: adam.core.spatial_math.ArrayLikeFactory

Generic factory. Give it (a) a Like class and (b) an xp namespace (array_api_compat.* if available; otherwise the library module).

_like[source]#
_xp[source]#
_dtype = None[source]#
_device = None[source]#
zeros(*shape) → ArrayAPILike[source]#
Parameters:

x (npt.ArrayLike) – matrix dimension

Returns:

zero matrix of dimension x

Return type:

npt.ArrayLike

eye(*shape) → ArrayAPILike[source]#
Parameters:

x (npt.ArrayLike) – matrix dimension

Returns:

identity matrix of dimension x

Return type:

npt.ArrayLike

asarray(x) → ArrayAPILike[source]#
Parameters:

x (npt.ArrayLike) – array

Returns:

array

Return type:

npt.ArrayLike

zeros_like(x: ArrayAPILike) → ArrayAPILike[source]#
Parameters:

x (npt.ArrayLike) – array

Returns:

one array with the same shape as x

Return type:

npt.ArrayLike

ones_like(x: ArrayAPILike) → ArrayAPILike[source]#
Parameters:

x (npt.ArrayLike) – array

Returns:

one array with the same shape as x

Return type:

npt.ArrayLike

tile(x: ArrayAPILike, reps: tuple) → ArrayAPILike[source]#
Parameters:

  • x (npt.ArrayLike) – input array

  • reps (tuple) – repetition factors for each dimension

Returns:

tiled array

Return type:

npt.ArrayLike

class adam.core.array_api_math.ArrayAPISpatialMath(factory, xp_getter: Callable[Ellipsis, Any] = xp_getter)[source]#

Bases: adam.core.spatial_math.SpatialMath

A drop-in SpatialMath that implements sin/cos/outer/concat/skew with the Array API.

Works for NumPy, PyTorch, and JAX; CasADi should keep its own subclass.

_xp_getter[source]#
_xp(*xs: Any)[source]#
sin(x)[source]#
Parameters:

x (npt.ArrayLike) – angle value

Returns:

sin value of x

Return type:

npt.ArrayLike

cos(x)[source]#
Parameters:

x (npt.ArrayLike) – angle value

Returns:

cos value of angle x

Return type:

npt.ArrayLike

skew(x)[source]#
outer(x, y)[source]#
vertcat(*x)[source]#
Parameters:

x (npt.ArrayLike) – elements

Returns:

vertical concatenation of elements x

Return type:

npt.ArrayLike

horzcat(*x)[source]#
Parameters:

x (npt.ArrayLike) – elements

Returns:

horizontal concatenation of elements x

Return type:

npt.ArrayLike

stack(x, axis=0)[source]#
Parameters:

  • x (npt.ArrayLike) – elements

  • axis (int) – axis along which to stack

Returns:

stacked elements x along axis

Return type:

npt.ArrayLike

concatenate(x, axis=0)[source]#
Parameters:

  • x (npt.ArrayLike) – elements

  • axis (int) – axis along which to concatenate

Returns:

concatenation of elements x along axis

Return type:

npt.ArrayLike

swapaxes(x: ArrayAPILike, axis1: int, axis2: int) → ArrayAPILike[source]#
expand_dims(x: ArrayAPILike, axis: int) → ArrayAPILike[source]#
transpose(x: ArrayAPILike, dims: tuple) → ArrayAPILike[source]#
Parameters:

  • x (npt.ArrayLike) – input array

  • dims (tuple) – permutation of dimensions

Returns:

transposed array

Return type:

npt.ArrayLike

inv(x: ArrayAPILike) → ArrayAPILike[source]#
Parameters:

x (npt.ArrayLike) – input array

Returns:

inverse of the array

Return type:

npt.ArrayLike

mtimes(A: ArrayAPILike, B: ArrayAPILike) → ArrayAPILike[source]#
solve(A: ArrayAPILike, B: ArrayAPILike) → ArrayAPILike[source]#
Contents