Flatten & Unflatten

Reshaping and axis manipulation layers.

class taktiny.nn.Flatten(start_axis=1, end_axis=-1)[source]

Bases: Module

Merge an inclusive, contiguous range of axes in row-major order.

Parameters:

  • start_axis (int) – First axis to merge; defaults to 1 to preserve a leading batch axis. Negative axes count from the end of the input.

  • end_axis (int) – Last axis to merge, inclusive; defaults to -1.

Axes outside the selected range, element order, and dtype are preserved. No batch or channel axes are inferred. A scalar can be flattened to (1,) using start_axis=0 (or -1); the default start_axis=1 requires rank >= 2. Zero-sized dimensions are supported. Axis bounds and order are checked against the input rank at call time. The module has no parameters or RNGs.

__call__ accepts out_sharding describing the output’s axes, not the input’s. With None, layout follows JAX’s reshape rules; it does not force replication.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> nn.Flatten()(jnp.ones((2, 3, 4))).shape
(2, 12)
>>> nn.Flatten(1, -2)(jnp.ones((2, 3, 4, 5))).shape
(2, 12, 5)
class taktiny.nn.Unflatten(axis, unflattened_size)[source]

Bases: Module

Replace one axis with a specified shape in row-major order.

Parameters:

  • axis (int) – Axis to expand. Negative axes count from the end of the input.

  • unflattened_size (int | Sequence[int]) – An integer or nonempty sequence of integer sizes. Sizes may be nonnegative, with at most one -1 for inference. The sequence is stored as an immutable tuple.

The new sizes must multiply to the selected axis size, independently of other axes (even if those axes have size zero). A -1 size is inferred when the product of known sizes divides the selected size. Combining -1 with zero is ambiguous and raises ValueError; inferring zero from a zero-sized axis and positive known sizes is supported. Scalar inputs are unsupported. Element order, dtype, and all other axes are preserved; no parameters or RNGs are used.

__call__ accepts out_sharding describing the expanded output axes. With None, JAX infers the layout; explicitly sharded inputs may require a target layout when splitting an axis has multiple possible sharding assignments.

Example

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> nn.Unflatten(-1, [3, -1])(jnp.ones((2, 12))).shape
(2, 3, 4)
>>> nn.Unflatten(0, (2, 0))(jnp.empty((0,))).shape
(2, 0)