Base Classes¶

Core abstractions for neural network modules and parameters in Taktiny.

class taktiny.nn.Module[source]¶

Bases: object

Base class for all neural network modules.

train()[source]¶

Sets the module and all its children to training mode.

Returns:: The module itself.
Return type:: Self

eval()[source]¶

Sets the module and all its children to evaluation mode.

Returns:: The module itself.
Return type:: Self

state_dict()[source]¶

Returns a hierarchical dictionary containing the module’s state.

Returns:: A nested dictionary representing the module state.
Return type:: dict[str, Any]

load_state_dict(state)[source]¶

Loads state values from a hierarchical dictionary into the module.

Parameters:: state (Mapping[str, Any]) – Hierarchical dictionary of state values.
Return type:: None

class taktiny.nn.Parameter(array, *, trainable=True, axis_names=None, partition_spec=None, metadata=None)[source]¶

Bases: Module

A PyTree node that wraps a JAX array or quantized tensor as a layer parameter.

Parameters delegate attribute access and arithmetic operations to their underlying values, allowing them to be used seamlessly in JAX mathematical operations.

Parameters:

array (Any) – The underlying tensor data (e.g., a jax.Array or qwix.QArray).
trainable (bool) – A metadata flag indicating if the parameter should be updated. Note that this does not automatically freeze the parameter in raw JAX; it must be explicitly filtered (e.g., by Taktiny’s Trainer or JAX tree utilities) before being passed to an optimizer. Defaults to True.
axis_names (tuple[str | None, ...] | None) – Logical axis names for advanced sharding or tensor parallelism. Defaults to None.
partition_spec (P | None) – Explicit hardware sharding specification. If provided along with axis_names, active logical mapping rules will override this value. Defaults to None.
metadata (dict[str, Any] | Sequence[tuple[str, Any]] | None) – Optional metadata dictionary for custom layer logic. Defaults to None.

Example

>>> import jax, jax.numpy as jnp
>>> from taktiny.nn import Parameter
>>>
>>> k = jax.random.key(0)
>>> k1, k2 = jax.random.split(k, 2)
>>> x = jax.random.normal(k1, (5, 10))
>>> z = Parameter(jax.random.normal(k2, (10, 5)))
>>> output = jnp.dot(x, z)

class taktiny.nn.Rngs(key, *, impl=None, dtype=None)[source]¶

Bases: object

A sequential random number generator class for maintaining PRNG state in JAX.

Parameters:

key (Union[Array, ndarray, bool, number, bool, int, float, complex]) – Seed or PRNGKey to initialize the state.
impl (Union[str, PRNGSpec, PRNGImpl, Hashable, None]) – PRNG implementation specification. Defaults to None.
dtype (Union[str, KeyTy, None]) – The dtype of the key array. Defaults to None.