traveler.core.typing.TypeMaker#

class traveler.core.typing.TypeMaker(shape, array_type=jt.Array)[source]#

Bases: object

A class to create a type for JAX arrays with a specific shape and dtype.

Parameters:

shape (tuple[int, ...]) – The shape of the JAX array.
dtype (str) – The data type of the JAX array.

Returns:

A JAX array type with the specified shape and dtype.

Return type:

jt.Array

Attributes#

`Categorical`
`Int8`
`Int16`
`Int32`
`Int64`
`Float32`
`Float64`
`Bool`