Arrays¶

tinax.array owns NumPy, JAX, DLPack, host materialization, logical array inspection, and a small set of validated array operations. Its conversion functions make copying and ownership explicit; its operations close silent-failure gaps in the equivalent jax.numpy/jax.nn calls.

See the Array Ownership guide for task-oriented usage.

from_numpy ¶

from_numpy(
    array: ndarray,
    *,
    copy: bool,
    device: Device | Sharding | None = None,
) -> Array

Create a JAX array from a NumPy array with explicit copy and placement policy.

Parameters:

Name	Type	Description	Default
`array`	`ndarray`	Source NumPy array to convert.	required
`copy`	`bool`	If `True`, take an independent host snapshot before transfer so that later mutation of `array` cannot affect the result. If `False`, allow JAX to alias the source buffer when the backend permits it.	required
`device`	`Device \| Sharding \| None`	Device or sharding to place the result on. `None` uses JAX's default placement.	`None`

Returns:

Type	Description
`Array`	A JAX array holding the converted data, placed according to `device`.

Raises:

Type	Description
`TypeError`	If `array` is not a `numpy.ndarray`, `copy` is not a `bool`, or `device` is not a `jax.Device`, `jax.sharding.Sharding`, or `None`.

from_dlpack ¶

from_dlpack(
    array: object,
    *,
    copy: bool | None = None,
    device: Device | Sharding | None = None,
) -> Array

Import a DLPack provider, completing copy=True before returning independent storage.

Parameters:

Name	Type	Description	Default
`array`	`object`	Object exposing the DLPack protocol to import.	required
`copy`	`bool \| None`	Copy policy. `True` forces an independent copy and blocks until the transfer completes; `False` requests a zero-copy import; `None` lets the backend choose.	`None`
`device`	`Device \| Sharding \| None`	Device or sharding for the imported array. `None` uses JAX's default placement.	`None`

Returns:

Type	Description
`Array`	A JAX array backed by the imported buffer. When `copy` is `True` the
`Array`	storage is guaranteed independent, including for same-device JAX providers.

Raises:

Type	Description
`TypeError`	If `copy` is not a `bool` or `None`, or `device` is not a `jax.Device`, `jax.sharding.Sharding`, or `None`.

to_numpy ¶

to_numpy(array: Array, *, writable: bool) -> ndarray

Synchronize a fully addressable JAX array and return a writable copy or read-only host value.

Blocks until the array is ready, then materializes its data on the host.

Parameters:

Name	Type	Description	Default
`array`	`Array`	Fully addressable JAX array to materialize on the host.	required
`writable`	`bool`	If `True`, return an independent, writable copy. If `False`, return a read-only NumPy value whose `writeable` flag is cleared.	required

Returns:

Type	Description
`ndarray`	The array's data as a `numpy.ndarray`. Writable results are independent
`ndarray`	copies; read-only results are a synchronized host view.

Raises:

Type	Description
`TypeError`	If `array` is not a `jax.Array` or `writable` is not a `bool`.
`ValueError`	If `array` is not fully addressable on the current process.

inspect_array ¶

inspect_array(array: Array) -> ArrayInfo

Inspect logical metadata without synchronizing or materializing array data.

Parameters:

Name	Type	Description	Default
`array`	`Array`	JAX array to inspect. Its data is never copied to the host.	required

Returns:

Type	Description
`ArrayInfo`	An `ArrayInfo` describing the array's shape, dtype, logical byte count,
`ArrayInfo`	commitment, and addressability.

Raises:

Type	Description
`TypeError`	If `array` is not a `jax.Array`.

ArrayInfo `dataclass` ¶

ArrayInfo(
    shape: tuple[int, ...],
    dtype: object,
    logical_nbytes: int | None,
    committed: bool,
    fully_addressable: bool,
)

Immutable logical metadata whose byte count is nominal rather than allocator memory.

Attributes:

Name	Type	Description
`shape`	`tuple[int, ...]`	Logical shape of the array.
`dtype`	`object`	Array dtype. Extended dtypes report `logical_nbytes` as `None`.
`logical_nbytes`	`int \| None`	Logical size in bytes, or `None` for extended dtypes whose element size is not a fixed byte count.
`committed`	`bool`	Whether the array is committed to a specific device.
`fully_addressable`	`bool`	Whether every shard is addressable by the current process.

stack_batch ¶

stack_batch(
    arrays: Sequence[Array], *, axis: int = 0
) -> Array

Stack arrays after requiring one shared dtype, where jnp.stack silently promotes mismatches.

Parameters:

Name	Type	Description	Default
`arrays`	`Sequence[Array]`	Non-empty sequence of JAX arrays to stack. All elements must share one dtype.	required
`axis`	`int`	Position of the new stacked axis. Defaults to 0.	`0`

Returns:

Type	Description
`Array`	The arrays stacked along `axis` via `jax.numpy.stack`.

Raises:

Type	Description
`TypeError`	If `arrays` is not a sequence, an element is not a `jax.Array`, or `axis` is not an integer (booleans rejected).
`ValueError`	If `arrays` is empty or its elements do not share one dtype.

safe_astype ¶

safe_astype(
    array: Array,
    dtype: DTypeLike,
    *,
    allow_lossy: bool = False,
) -> Array

Cast to dtype, requiring allow_lossy=True for a narrowing cast that jnp.astype performs silently.

Parameters:

Name	Type	Description	Default
`array`	`Array`	JAX array to cast.	required
`dtype`	`DTypeLike`	Target dtype.	required
`allow_lossy`	`bool`	Whether to permit a cast that can lose precision or overflow. Defaults to `False`.	`False`

Returns:

Type	Description
`Array`	`array` cast to `dtype` via `jax.numpy.astype`.

Raises:

Type	Description
`TypeError`	If `array` is not a `jax.Array` or `allow_lossy` is not a bool.
`ValueError`	If the cast from `array`'s dtype to `dtype` is not value-preserving and `allow_lossy` is `False`.

one_hot ¶

one_hot(
    indices: Array,
    *,
    num_classes: int,
    dtype: DTypeLike | None = None,
) -> Array

Encode integer indices as one-hot rows, rejecting inputs jax.nn.one_hot would silently zero out.

Does not validate that concrete index values lie in 0..num_classes - 1: that requires concrete values and would break tracing under jax.jit. Out-of-range indices still produce an all-zero row, matching jax.nn.one_hot's own documented behavior.

Parameters:

Name	Type	Description	Default
`indices`	`Array`	JAX array of integer indices to encode.	required
`num_classes`	`int`	Number of one-hot classes. Must be a positive integer.	required
`dtype`	`DTypeLike \| None`	Output dtype. `None` uses `jax.nn.one_hot`'s default.	`None`

Returns:

Type	Description
`Array`	A one-hot encoding of `indices` with a trailing `num_classes` axis.

Raises:

Type	Description
`TypeError`	If `indices` is not a `jax.Array`, its dtype is not an integer dtype, or `num_classes` is not an integer (booleans rejected).
`ValueError`	If `num_classes` is not positive.

Arrays¶

from_numpy ¶

from_dlpack ¶

to_numpy ¶

inspect_array ¶

ArrayInfo dataclass ¶

stack_batch ¶

safe_astype ¶

one_hot ¶

ArrayInfo `dataclass` ¶