Sequence Mixer¤
linax.sequence_mixers.base.SequenceMixerConfig
¤
Configuration for sequence mixers.
Attributes:
| Name | Type | Description |
|---|---|---|
state_dim |
Dimensionality of the state space. |
build(in_features: int, key: PRNGKeyArray) -> SequenceMixer
¤
Build sequence mixer from config.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
in_features
|
int
|
Input dimensionality. |
required |
key
|
PRNGKeyArray
|
JAX random key for initialization. |
required |
Returns:
| Type | Description |
|---|---|
SequenceMixer
|
The sequence mixer instance. |
linax.sequence_mixers.base.SequenceMixer
¤
Abstract base class for all sequence mixers.
This class is used to define the interface for all sequence mixers.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
in_features
|
int
|
Input dimensionality. |
required |
cfg
|
ConfigType
|
Configuration for the sequence mixer. |
required |
key
|
PRNGKeyArray
|
JAX random key for initialization. |
required |
__init__(in_features: int, cfg: ConfigType, key: PRNGKeyArray)
¤
Initialize the sequence mixer.
__call__(x: Array, key: PRNGKeyArray) -> Array
¤
Forward pass of the sequence mixer.
Parameters:
| Name | Type | Description | Default |
|---|---|---|---|
x
|
Array
|
The input sequence to the sequence mixer. |
required |
key
|
PRNGKeyArray
|
The random key for the sequence mixer. |
required |
Returns:
| Type | Description |
|---|---|
Array
|
The output of the sequence mixer. |