LSTMParams

chronax.lstm.LSTMParams · inherits NamedTuple

Trainable parameters for one LSTM layer and a linear readout.

LSTMState

chronax.lstm.LSTMState · inherits NamedTuple

Carry state for an LSTM layer.

LSTMModel

chronax.lstm.LSTMModel

A pure-JAX LSTM followed by a linear projection.

Inputs are expected in time-major form for a single sequence, (sequence_length, input_size). Batched inputs use (batch_size, sequence_length, input_size) and are evaluated with :meth:apply_batch.

__init__(self, input_size, hidden_size, output_size)

Parameter	Type	Default	Description
input_size	int	-	(undocumented)
hidden_size	int	-	(undocumented)
output_size	int	-	(undocumented)

init(self, key: Array) -> LSTMParams

Initialize model parameters with Xavier-style uniform weights.

Parameter	Type	Default	Description
key	Array	-	(undocumented)

Returns: LSTMParams

initial_state(self) -> LSTMState

Return a zero-valued recurrent state.

Returns: LSTMState

step(self, params: LSTMParams, state: LSTMState, x_t: Array) -> LSTMState

Run one LSTM step.

Parameter	Type	Default	Description
params	LSTMParams	-	(undocumented)
state	LSTMState	-	(undocumented)
x_t	Array	-	(undocumented)

Returns: LSTMState

apply(self, params: LSTMParams, inputs: Array, state: LSTMState | None = None) -> tuple[Array, LSTMState]

Evaluate one sequence.

Parameter	Type	Default	Description
params	LSTMParams	-	(undocumented)
inputs	Array	-	(undocumented)
state	LSTMState | None	None	(undocumented)

Returns: tuple[Array, LSTMState] (A tuple of (outputs, final_state) where outputs has shape (sequence_length, output_size).)

apply_batch(self, params: LSTMParams, inputs: Array) -> tuple[Array, LSTMState]

Evaluate a batch of sequences.

Parameter	Type	Default	Description
params	LSTMParams	-	(undocumented)
inputs	Array	-	(undocumented)

Returns: tuple[Array, LSTMState] (The outputs and final state for the batch.)

xavier_uniform

chronax.lstm.xavier_uniform

Sample Xavier/Glorot uniform weights for a matrix-like shape.

Parameter	Type	Default	Description
key	Array	-	(undocumented)
shape	tuple[int, ...]	-	(undocumented)

Returns: Array Raises: ValueError