xlstm_jax.kernels.mlstm_recurrent.triton_fused_fw

Functions

recurrent_step_fw(matC_state, vecN_state, scaM_state, ...)

Module Contents

xlstm_jax.kernels.mlstm_recurrent.triton_fused_fw.recurrent_step_fw(matC_state, vecN_state, scaM_state, vecQ, vecK, vecV, scaI, scaF, matC_new=None, vecN_new=None, scaM_new=None, qk_scale=None, eps=1e-06)

Parameters:

• matC_state (jax.Array)

• vecN_state (jax.Array)

• scaM_state (jax.Array)

• vecQ (jax.Array)

• vecK (jax.Array)

• vecV (jax.Array)

• scaI (jax.Array)

• scaF (jax.Array)

• matC_new (jax.Array | None)

• vecN_new (jax.Array | None)

• scaM_new (jax.Array | None)

• qk_scale (float | None)

• eps (float)