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

flat_state_dict(prefix='', *, include=None)[source]¶

Returns a flattened dictionary containing the module’s state.

Parameters:

prefix (str) – A prefix to prepend to all keys. Defaults to ‘’.
include (Sequence[str] | None) – Full-match regexes for dotted parameter paths, including prefix. None includes all paths; an empty sequence includes none. Nonempty patterns matching nothing raise ValueError.

Returns:

A dictionary mapping flattened parameter names to their values.

Return type:

dict[str, Any]

flat_parameter_dict(prefix='')[source]¶

Returns a flattened dictionary containing the module’s parameters.

Parameters:: prefix (str) – A prefix to prepend to all keys. Defaults to ‘’.
Returns:: A dictionary mapping flattened names to Parameter objects.
Return type:: dict[str, Any]

state_dict(*, include=None)[source]¶

Returns a hierarchical dictionary containing the module’s state.

Parameters:: include (Sequence[str] | None) – Full-match regexes for dotted parameter paths. None includes all paths; an empty sequence includes none. Empty branches are omitted when filtering. Nonempty patterns matching nothing raise ValueError.
Returns:: A nested dictionary representing the module state.
Return type:: dict[str, Any]

load_flat_state_dict(state, prefix='', *, include=None)[source]¶

Loads state values from a flattened dictionary into the module.

Parameters:

state (Mapping[str, Any]) – Flattened dictionary of state values.
prefix (str) – Prefix used in the flattened keys. Defaults to ‘’.
include (Sequence[str] | None) – Full-match regexes for dotted target parameter paths, including prefix. Only selected keys present in state are loaded. None selects all; an empty sequence selects none. Nonempty patterns matching nothing raise ValueError.

Return type:

None

load_state_dict(state, *, include=None)[source]¶

Loads state values from a hierarchical dictionary into the module.

Parameters:

state (Mapping[str, Any]) – Hierarchical dictionary of state values.
include (Sequence[str] | None) – Full-match regexes for dotted target parameter paths. Only selected keys present in state are loaded. None selects all; an empty sequence selects none. Nonempty patterns matching nothing raise ValueError.

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)

property value: PyTree¶: Returns the underlying tensor data.

taktiny.nn.module(cls)[source]¶

Class decorator to transform a generic class into a Module subclass.

Returns:: The newly created Module subclass.
Return type:: _type_

class taktiny.nn.Pytree[source]¶: Bases: object

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.

split_key(num_splits)[source]¶

Return independent keys and advance this stream once.

Parameters:: num_splits (int) – Number of keys to return; must be a positive integer.
Return type:: tuple[Array, ...]
Returns:: A tuple of num_splits keys. The first key from the split is retained as this stream’s new state.

split_rngs(num_splits)[source]¶

Return independent Rngs streams and advance this stream once.

Parameters:: num_splits (int) – Number of streams to return; must be a positive integer.
Return type:: tuple[Rngs, ...]
Returns:: A tuple of num_splits streams, each initialized from a distinct key returned by split_key().

property key: PRNGKey¶

Return the current PRNGKey state.

Returns:: The current PRNG state.
Return type:: PRNGKey