Index

A

abs_squared_norm() (in module lgatr.primitives.invariants) additional_qk_mv_channels (lgatr.layers.attention.config.SelfAttentionConfig attribute) additional_qk_s_channels (lgatr.layers.attention.config.SelfAttentionConfig attribute)

attn_ratio (lgatr.layers.attention.config.CrossAttentionConfig attribute) (lgatr.layers.attention.config.SelfAttentionConfig attribute) autocast_dtype() (in module lgatr.utils.autocast)

B

bivector (lgatr.primitives.config.PrimitivesConfig attribute)

C

cast() (lgatr.layers.attention.config.CrossAttentionConfig class method) (lgatr.layers.attention.config.SelfAttentionConfig class method) (lgatr.layers.mlp.config.MLPConfig class method) (lgatr.primitives.config.PrimitivesConfig class method) (lgatr.utils.autocast.minimum_autocast_precision method) compile_model() (in module lgatr.utils.compile)

ConditionalLGATr (class in lgatr.nets.conditional_lgatr) ConditionalLGATrBlock (class in lgatr.layers.conditional_lgatr_block) ConditionalLGATrSlim (class in lgatr.nets.conditional_slim) ConditionalSlimBlock (class in lgatr.layers.slim_layers) CrossAttention (class in lgatr.layers.attention.cross_attention) CrossAttentionConfig (class in lgatr.layers.attention.config)

D

dropout_prob (lgatr.layers.attention.config.CrossAttentionConfig attribute) (lgatr.layers.attention.config.SelfAttentionConfig attribute) (lgatr.layers.mlp.config.MLPConfig attribute)

E

embed_axialvector() (in module lgatr.interface.axialvector) embed_bivector() (in module lgatr.interface.bivector) embed_pseudoscalar() (in module lgatr.interface.pseudoscalar) embed_scalar() (in module lgatr.interface.scalar) embed_vector() (in module lgatr.interface.vector) equi_layer_norm() (in module lgatr.primitives.normalization) equi_linear() (in module lgatr.primitives.linear)

EquiLayerNorm (class in lgatr.layers.layer_norm) EquiLinear (class in lgatr.layers.linear) extract_axialvector() (in module lgatr.interface.axialvector) extract_bivector() (in module lgatr.interface.bivector) extract_pseudoscalar() (in module lgatr.interface.pseudoscalar) extract_scalar() (in module lgatr.interface.scalar) extract_vector() (in module lgatr.interface.vector)

F

forward() (lgatr.layers.attention.cross_attention.CrossAttention method) (lgatr.layers.attention.self_attention.SelfAttention method) (lgatr.layers.conditional_lgatr_block.ConditionalLGATrBlock method) (lgatr.layers.dropout.GradeDropout method) (lgatr.layers.layer_norm.EquiLayerNorm method) (lgatr.layers.lgatr_block.LGATrBlock method) (lgatr.layers.linear.EquiLinear method) (lgatr.layers.mlp.geometric_bilinears.GeometricBilinear method) (lgatr.layers.mlp.mlp.GeoMLP method) (lgatr.layers.mlp.nonlinearities.ScalarGatedNonlinearity method) (lgatr.layers.slim_layers.ConditionalSlimBlock method) (lgatr.layers.slim_layers.SlimBlock method) (lgatr.layers.slim_layers.SlimCrossAttention method) (lgatr.layers.slim_layers.SlimDropout method) (lgatr.layers.slim_layers.SlimGLU method) (lgatr.layers.slim_layers.SlimLinear method) (lgatr.layers.slim_layers.SlimMLP method) (lgatr.layers.slim_layers.SlimRMSNorm method) (lgatr.layers.slim_layers.SlimSelfAttention method) (lgatr.nets.conditional_lgatr.ConditionalLGATr method) (lgatr.nets.conditional_slim.ConditionalLGATrSlim method) (lgatr.nets.lgatr.LGATr method) (lgatr.nets.slim.LGATrSlim method)

G

geometric_product (lgatr.primitives.config.PrimitivesConfig attribute) geometric_product() (in module lgatr.primitives.bilinear) GeometricBilinear (class in lgatr.layers.mlp.geometric_bilinears) GeoMLP (class in lgatr.layers.mlp.mlp) get_attention_backend() (in module lgatr.primitives.attention_backends)

get_num_spurions() (in module lgatr.interface.spurions) get_spurions() (in module lgatr.interface.spurions) grade_dropout() (in module lgatr.primitives.dropout) grade_involute() (in module lgatr.primitives.linear) grade_project() (in module lgatr.primitives.linear) GradeDropout (class in lgatr.layers.dropout)

H

head_scale (lgatr.layers.attention.config.CrossAttentionConfig attribute) (lgatr.layers.attention.config.SelfAttentionConfig attribute) hidden_mv_channels (lgatr.layers.attention.config.CrossAttentionConfig property) (lgatr.layers.attention.config.SelfAttentionConfig property)

hidden_s_channels (lgatr.layers.attention.config.CrossAttentionConfig property) (lgatr.layers.attention.config.SelfAttentionConfig property)

I

in_mv_channels (lgatr.layers.attention.config.SelfAttentionConfig attribute)

in_s_channels (lgatr.layers.attention.config.SelfAttentionConfig attribute) inner_product() (in module lgatr.primitives.invariants)

K

kv_mv_channels (lgatr.layers.attention.config.CrossAttentionConfig attribute)

kv_s_channels (lgatr.layers.attention.config.CrossAttentionConfig attribute)

L

LGATr (class in lgatr.nets.lgatr) lgatr.interface.axialvector module lgatr.interface.bivector module lgatr.interface.pseudoscalar module lgatr.interface.scalar module lgatr.interface.spurions module lgatr.interface.vector module lgatr.primitives.attention module lgatr.primitives.bilinear module

lgatr.primitives.compile module lgatr.primitives.dropout module lgatr.primitives.invariants module lgatr.primitives.linear module lgatr.primitives.normalization module lgatr.utils.autocast module lgatr.utils.compile module LGATrBlock (class in lgatr.layers.lgatr_block) LGATrSlim (class in lgatr.nets.slim)

M

minimum_autocast_precision (class in lgatr.utils.autocast) mlp_ratio (lgatr.layers.mlp.config.MLPConfig attribute) MLPConfig (class in lgatr.layers.mlp.config) module lgatr.interface.axialvector lgatr.interface.bivector lgatr.interface.pseudoscalar lgatr.interface.scalar lgatr.interface.spurions lgatr.interface.vector lgatr.primitives.attention lgatr.primitives.bilinear lgatr.primitives.compile lgatr.primitives.dropout lgatr.primitives.invariants lgatr.primitives.linear lgatr.primitives.normalization lgatr.utils.autocast lgatr.utils.compile

multi_query (lgatr.layers.attention.config.CrossAttentionConfig attribute) (lgatr.layers.attention.config.SelfAttentionConfig attribute) mv_channels (lgatr.layers.mlp.config.MLPConfig attribute)

N

naive_amp (class in lgatr.utils.autocast) nonlinearity (lgatr.layers.mlp.config.MLPConfig attribute) num_heads (lgatr.layers.attention.config.CrossAttentionConfig attribute) (lgatr.layers.attention.config.SelfAttentionConfig attribute)

num_layers_mlp (lgatr.layers.mlp.config.MLPConfig attribute) num_pin_linear_basis_elements (lgatr.primitives.config.PrimitivesConfig property)

O

out_mv_channels (lgatr.layers.attention.config.CrossAttentionConfig attribute) (lgatr.layers.attention.config.SelfAttentionConfig attribute) out_s_channels (lgatr.layers.attention.config.CrossAttentionConfig attribute) (lgatr.layers.attention.config.SelfAttentionConfig attribute)

output_init (lgatr.layers.attention.config.CrossAttentionConfig attribute) (lgatr.layers.attention.config.SelfAttentionConfig attribute)

P

PrimitivesConfig (class in lgatr.primitives.config)

Q

q_mv_channels (lgatr.layers.attention.config.CrossAttentionConfig attribute)

q_s_channels (lgatr.layers.attention.config.CrossAttentionConfig attribute)

R

reset_parameters() (lgatr.layers.linear.EquiLinear method) (lgatr.layers.slim_layers.SlimLinear method)

reverse() (in module lgatr.primitives.linear)

S

s_channels (lgatr.layers.mlp.config.MLPConfig attribute) ScalarGatedNonlinearity (class in lgatr.layers.mlp.nonlinearities) scaled_dot_product_attention() (in module lgatr.primitives.attention) sdp_attention() (in module lgatr.primitives.attention) SelfAttention (class in lgatr.layers.attention.self_attention) SelfAttentionConfig (class in lgatr.layers.attention.config) SlimBlock (class in lgatr.layers.slim_layers) SlimCrossAttention (class in lgatr.layers.slim_layers)

SlimDropout (class in lgatr.layers.slim_layers) SlimGLU (class in lgatr.layers.slim_layers) SlimLinear (class in lgatr.layers.slim_layers) SlimMLP (class in lgatr.layers.slim_layers) SlimRMSNorm (class in lgatr.layers.slim_layers) SlimSelfAttention (class in lgatr.layers.slim_layers) sparse_gp (lgatr.primitives.config.PrimitivesConfig attribute) sparse_linear (lgatr.primitives.config.PrimitivesConfig attribute) subgroup (lgatr.primitives.config.PrimitivesConfig attribute)

W

warmup_after_apply() (in module lgatr.utils.compile)

warmup_caches() (in module lgatr.primitives.compile)