Skip to content

Identity Sequence Mixer¤

discretax.sequence_mixers.identity.IdentitySequenceMixer ¤

Identity sequence mixer layer.

This layer implements a simple identity/pass-through operation that returns the input sequence unchanged.

Parameters:

Name Type Description Default
in_features int

Input dimensionality.

required
key PRNGKeyArray

JAX random key for initialization.

required
__init__(in_features: int, key: PRNGKeyArray, *args, **kwargs) ¤

Initialize the Identity sequence mixer layer.

__call__(x: Array, key: PRNGKeyArray) -> Array ¤

Forward pass of the Identity sequence mixer layer.

Parameters:

Name Type Description Default
x Array

Input sequence of features.

required
key PRNGKeyArray

JAX random key (unused, for compatibility).

required

Returns:

Type Description
Array

The input sequence unchanged.