Index

_

__init__() (CentaurusBase method) (Classifier method) (ClassifierBlock method) (FFTConvS4 method) (LRNN method) (LRU method), [1] (LTI_LRNN method), [1] (LTV_LRNN method), [1] (Mamba method), [1] (RGLRU method), [1] (S4 method), [1] (S4D method), [1] (S5 method), [1], [2], [3] (S7 method), [1] (SequencePooling method)

A

allocate_inference_cache() (CentaurusBase method) (CentaurusPWNeck method), [1] (LRNNLMHeadModel method) (LRNNModel method) (LRU method), [1] (LTI_LRNN method), [1] (LTV_LRNN method), [1] (Mamba method), [1] (RGLRU method), [1] (S4 method), [1] (S4D method), [1] (S5 method), [1], [2], [3] (S7 method), [1]

async_() (in module lrnnx.core.discretization)

B

backward() (LayerNormFn static method) (LayerNormLinearFn static method) (MambaInnerFn static method) (RGLRUInnerFn static method) (RGLRUScanFn static method) (S5InnerFn static method) (S7InnerFn static method) (S7ScanFn static method) (SelectiveScanFn static method) (SimplifiedScanFn static method)

bilinear() (in module lrnnx.core.discretization)

C

capture_graph() (in module lrnnx.utils.generation) cauchy_naive() (in module lrnnx.ops.s4_utils) Centaurus (class in lrnnx.models.lti) (class in lrnnx.models.lti.centaurus) CentaurusBase (class in lrnnx.models.lti.centaurus) CentaurusDWS (class in lrnnx.models.lti) (class in lrnnx.models.lti.centaurus) CentaurusFull (class in lrnnx.models.lti) (class in lrnnx.models.lti.centaurus) CentaurusNeck (class in lrnnx.models.lti) (class in lrnnx.models.lti.centaurus) CentaurusPWNeck (class in lrnnx.models.lti) (class in lrnnx.models.lti.centaurus)

Classifier (class in lrnnx.architectures.classifier) ClassifierBlock (class in lrnnx.architectures.classifier) combination() (in module lrnnx.ops.s4_utils) compute_kernel() (CentaurusBase method) (CentaurusPWNeck method), [1] (LRU method), [1] (LTI_LRNN method), [1] (S5 method), [1] config_prune() (in module lrnnx.ops.triton.layer_norm) create_block() (in module lrnnx.architectures.language_model) CUDAGraphStepCache (class in lrnnx.utils.generation) custom_amp_decorator() (in module lrnnx.ops.torch) custom_bwd() (in module lrnnx.ops.torch) custom_fwd() (in module lrnnx.ops.torch)

D

d_output (FFTConvS4 property) (S4 property), [1] (S4D property), [1] default_state() (FFTConvS4 method) (S4 method), [1] (S4D method), [1] (S4DKernel method) (S4Kernel method)

dirac() (in module lrnnx.core.discretization) discretize() (CentaurusBase method) (LRU method), [1] (LTI_LRNN method), [1] (S5 method), [1] double_length() (S4Kernel method) DownPool1D (class in lrnnx.architectures.lru_unet) dplr() (in module lrnnx.ops.s4_utils) DropoutNd (class in lrnnx.ops.s4_utils)

F

fft_conv() (in module lrnnx.core.convolution) FFTConvS4 (class in lrnnx.core.convolution) forward() (CentaurusBase method) (CentaurusPWNeck method), [1] (Classifier method) (ClassifierBlock method) (DownPool1D method) (DropoutNd method) (FFTConvS4 method) (LayerNormFeature method) (LayerNormFn static method) (LayerNormLinearFn static method) (LRNN method) (LRNNLMHeadModel method) (LRNNModel method) (LRU method), [1] (LRU_UNet method) (LTV_LRNN method), [1] (Mamba method), [1] (MambaInnerFn static method) (PositionEmbedding method) (RGLRU method), [1] (RGLRUInnerFn static method) (RGLRUScanFn static method) (RMSNorm method) (S4 method), [1] (S4D method), [1] (S4DKernel method) (S4Kernel method) (S5 method), [1], [2], [3] (S5InnerFn static method) (S7 method), [1] (S7InnerFn static method) (S7ScanFn static method) (SelectiveScanFn static method) (SequencePooling method) (SimplifiedScanFn static method) (TokenEmbedding method) (UpPool1D method)

forward_state() (S4DKernel method) (S4Kernel method) from_pretrained() (LRNNLMHeadModel class method)

G

generate() (in module lrnnx.utils.generation) get_cauchy_kernel() (in module lrnnx.ops.s4_utils)

get_vandermonde_kernel() (in module lrnnx.ops.s4_utils) get_vandermonde_transpose_kernel() (in module lrnnx.ops.s4_utils)

I

init_CV() (in module lrnnx.utils.init) init_dt() (in module lrnnx.ops.s4_utils) init_log_steps() (in module lrnnx.utils.init)

init_ssm_dplr() (in module lrnnx.ops.s4_utils) init_VinvB() (in module lrnnx.utils.init) inv_transform() (in module lrnnx.ops.s4_utils)

L

layer_norm_fn() (in module lrnnx.ops.triton.layer_norm) layer_norm_linear_fn() (in module lrnnx.ops.triton.layer_norm) layer_norm_ref() (in module lrnnx.ops.triton.layer_norm) LayerNormFeature (class in lrnnx.architectures.lru_unet) LayerNormFn (class in lrnnx.ops.triton.layer_norm) LayerNormLinearFn (class in lrnnx.ops.triton.layer_norm) LinearActivation() (in module lrnnx.ops.s4_utils) log_vandermonde_naive() (in module lrnnx.ops.s4_utils) log_vandermonde_transpose_naive() (in module lrnnx.ops.s4_utils) LRNN (class in lrnnx.core.base) LRNNLMHeadModel (class in lrnnx.architectures.language_model) LRNNModel (class in lrnnx.architectures.language_model) lrnnx.architectures.classifier module lrnnx.architectures.embedding module lrnnx.architectures.language_model module lrnnx.architectures.lru_unet module lrnnx.core.base module lrnnx.core.convolution module lrnnx.core.discretization module lrnnx.models.lti module lrnnx.models.lti.base module lrnnx.models.lti.centaurus module lrnnx.models.lti.lru module lrnnx.models.lti.s4 module lrnnx.models.lti.s4d module lrnnx.models.lti.s5 module lrnnx.models.ltv module lrnnx.models.ltv.base module

lrnnx.models.ltv.mamba module lrnnx.models.ltv.rglru module lrnnx.models.ltv.s5 module lrnnx.models.ltv.s7 module lrnnx.ops.rglru_scan module lrnnx.ops.s4_kernel_interface module lrnnx.ops.s4_utils module lrnnx.ops.s7_scan module lrnnx.ops.selective_scan module lrnnx.ops.simplified_scan module lrnnx.ops.torch module lrnnx.ops.triton module lrnnx.ops.triton.layer_norm module lrnnx.ops.triton.selective_state_update module lrnnx.ops.triton.simplified_state_update module lrnnx.ops.triton.softplus module lrnnx.utils.generation module lrnnx.utils.init module LRU (class in lrnnx.models.lti) (class in lrnnx.models.lti.lru) LRU_UNet (class in lrnnx.architectures.lru_unet) LTI_LRNN (class in lrnnx.models.lti) (class in lrnnx.models.lti.base) LTV_LRNN (class in lrnnx.models.ltv) (class in lrnnx.models.ltv.base)

M

make_DPLR_HiPPO() (in module lrnnx.utils.init) make_HiPPO() (in module lrnnx.utils.init) make_NPLR_HiPPO() (in module lrnnx.utils.init) Mamba (class in lrnnx.models.ltv) (class in lrnnx.models.ltv.mamba) mamba_inner_fn() (in module lrnnx.ops.selective_scan) mamba_inner_ref() (in module lrnnx.ops.selective_scan) MambaInnerFn (class in lrnnx.ops.selective_scan) module lrnnx.architectures.classifier lrnnx.architectures.embedding lrnnx.architectures.language_model lrnnx.architectures.lru_unet lrnnx.core.base lrnnx.core.convolution lrnnx.core.discretization lrnnx.models.lti lrnnx.models.lti.base lrnnx.models.lti.centaurus lrnnx.models.lti.lru lrnnx.models.lti.s4 lrnnx.models.lti.s4d lrnnx.models.lti.s5 lrnnx.models.ltv lrnnx.models.ltv.base lrnnx.models.ltv.mamba lrnnx.models.ltv.rglru lrnnx.models.ltv.s5 lrnnx.models.ltv.s7 lrnnx.ops.rglru_scan lrnnx.ops.s4_kernel_interface lrnnx.ops.s4_utils lrnnx.ops.s7_scan lrnnx.ops.selective_scan lrnnx.ops.simplified_scan lrnnx.ops.torch lrnnx.ops.triton lrnnx.ops.triton.layer_norm lrnnx.ops.triton.selective_state_update lrnnx.ops.triton.simplified_state_update lrnnx.ops.triton.softplus lrnnx.utils.generation lrnnx.utils.init

N

no_discretization() (in module lrnnx.core.discretization)

nplr() (in module lrnnx.ops.s4_utils)

O

opt_ssm_forward() (in module lrnnx.core.convolution)

P

param_transform() (in module lrnnx.ops.s4_utils) PositionEmbedding (class in lrnnx.architectures.embedding)

power() (in module lrnnx.ops.s4_utils) process_dplr_params() (in module lrnnx.ops.s4_utils) process_ssm_params() (in module lrnnx.ops.s4_utils)

R

rank_correction() (in module lrnnx.ops.s4_utils) register_ssm_params() (in module lrnnx.ops.s4_utils) reset_parameters() (RMSNorm method) RGLRU (class in lrnnx.models.ltv) (class in lrnnx.models.ltv.rglru) rglru_inner_fn() (in module lrnnx.ops.rglru_scan) rglru_inner_ref() (in module lrnnx.ops.rglru_scan)

rglru_scan_fn() (in module lrnnx.ops.rglru_scan) rglru_scan_ref() (in module lrnnx.ops.rglru_scan) RGLRUInnerFn (class in lrnnx.ops.rglru_scan) RGLRUScanFn (class in lrnnx.ops.rglru_scan) rms_norm_fn() (in module lrnnx.ops.triton.layer_norm) rms_norm_forward() (in module lrnnx.ops.selective_scan) rms_norm_ref() (in module lrnnx.ops.triton.layer_norm) RMSNorm (class in lrnnx.ops.triton.layer_norm)

S

S4 (class in lrnnx.models.lti) (class in lrnnx.models.lti.s4) S4D (class in lrnnx.models.lti) (class in lrnnx.models.lti.s4d) S4DKernel (class in lrnnx.ops.s4_kernel_interface) S4Kernel (class in lrnnx.ops.s4_kernel_interface) S4KernelBase (class in lrnnx.ops.s4_kernel_interface) S5 (class in lrnnx.models.lti) (class in lrnnx.models.lti.s5) (class in lrnnx.models.ltv) (class in lrnnx.models.ltv.s5) s5_inner_fn() (in module lrnnx.ops.simplified_scan) s5_inner_ref() (in module lrnnx.ops.simplified_scan) S5InnerFn (class in lrnnx.ops.simplified_scan) S7 (class in lrnnx.models.ltv) (class in lrnnx.models.ltv.s7) s7_inner_fn() (in module lrnnx.ops.s7_scan) s7_inner_ref() (in module lrnnx.ops.s7_scan) s7_scan_fn() (in module lrnnx.ops.s7_scan) s7_scan_ref() (in module lrnnx.ops.s7_scan) S7InnerFn (class in lrnnx.ops.s7_scan) S7ScanFn (class in lrnnx.ops.s7_scan) save_pretrained() (LRNNLMHeadModel method) selective_scan_fn() (in module lrnnx.ops.selective_scan) selective_scan_ref() (in module lrnnx.ops.selective_scan) selective_state_update() (in module lrnnx.ops.triton.selective_state_update) selective_state_update_ref() (in module lrnnx.ops.triton.selective_state_update)

SelectiveScanFn (class in lrnnx.ops.selective_scan) SequencePooling (class in lrnnx.architectures.classifier) setup_default_state() (in module lrnnx.ops.s4_utils) setup_step() (FFTConvS4 method) simplified_scan_fn() (in module lrnnx.ops.simplified_scan) simplified_scan_ref() (in module lrnnx.ops.simplified_scan) simplified_state_update() (in module lrnnx.ops.triton.simplified_state_update) simplified_state_update_ref() (in module lrnnx.ops.triton.simplified_state_update) SimplifiedScanFn (class in lrnnx.ops.simplified_scan) softplus() (in module lrnnx.ops.triton.softplus) ssm() (in module lrnnx.ops.s4_utils) step() (CentaurusBase method) (CentaurusPWNeck method), [1] (FFTConvS4 method) (LRNNLMHeadModel method) (LRNNModel method) (LRU method), [1] (LTI_LRNN method), [1] (LTV_LRNN method), [1] (Mamba method), [1] (RGLRU method), [1] (S4 method), [1] (S4D method), [1] (S4DKernel method) (S4Kernel method) (S5 method), [1], [2], [3] (S7 method), [1]

T

tie_weights() (LRNNLMHeadModel method)

TokenEmbedding (class in lrnnx.architectures.embedding) transition() (in module lrnnx.ops.s4_utils)

U

UpPool1D (class in lrnnx.architectures.lru_unet)

Z

zoh() (in module lrnnx.core.discretization)