typinox.shaped module

typinox.shaped module#

class typinox.shaped.NDArray(*args, **kwargs)[source]#

Bases: Protocol

Duck match for any object that behaves like a numpy array (has .shape and .dtype attributes).

typinox.shaped.ensure_shape(shape: int | tuple[int, ...] | NDArray, dim_spec: str, /)[source]#
typinox.shaped.ensure_shape(name: str, shape: int | tuple[int, ...] | NDArray, dim_spec: str, /)

Ensure that the shape of an array matches jaxtyping named dimensions.

Parameters:
  • name (str, optional) – The name of the array.

  • shape (ShapeLike) – The shape of the array.

  • dim_spec (str) – The jaxtyping named dimensions to be matched.

Returns:

If the shape matches the named dimensions, return None.

Return type:

None

Raises:

ValidationFailed – when the shape does not match the named dimensions

typinox.shaped.ensure_shape_equal(shape1: int | tuple[int, ...] | NDArray, shape2: int | tuple[int, ...] | NDArray, /)[source]#
typinox.shaped.ensure_shape_equal(name: str, shape1: int | tuple[int, ...] | NDArray, shape2: int | tuple[int, ...] | NDArray, /)
typinox.shaped.ensure_shape_equal(name1: str, shape1: int | tuple[int, ...] | NDArray, name2: str, shape2: int | tuple[int, ...] | NDArray, /)

Ensure that the shapes of two arrays are equal.

Parameters:
  • name (str, optional) – The name of the arrays. Cannot be provided if name1 or name2 is provided.

  • name1 (str, optional) – The name of the first array. Must be provided if name2 is provided.

  • shape1 (int, or tuple[int, …], or an object with .shape) – The shape of the first array.

  • name2 (str, optional) – The name of the second array. Must be provided if name1 is provided.

  • shape2 (int, or tuple[int, …], or an object with .shape) – The shape of the second array.

Returns:

If the shapes are equal, return None.

Return type:

None

Raises:

ValidationFailed – when the shapes are not equal.