Index _ | A | B | C | D | E | F | G | H | I | K | L | M | O | P | Q | R | S | T | U | V | W _ __init__() (taktiny.data.DataLoader method) _apply() (taktiny.nn.AlphaDropout method) (taktiny.nn.Dropout method) _scaled_size() (taktiny.nn.Downsample method) (taktiny.nn.Upsample method) _spatial_shape() (taktiny.nn.Downsample method) _transpose_padding() (taktiny.nn.ConvTranspose static method) A ActivationBase (class in taktiny.nn) AdaLoRAAdapter (class in taktiny.takt) AdaLoRALinear (class in taktiny.nn) AdaptiveAvgPool (class in taktiny.nn) AdaptiveMaxPool (class in taktiny.nn) add_callback() (taktiny.trainer.Trainer method) AlphaDropout (class in taktiny.nn) apply_adapter() (taktiny.takt.Takt class method) ApplyTemplate (class in taktiny.data) Attention (class in taktiny.nn) AvgPool (class in taktiny.nn) B BaseAdapter (in module taktiny.takt) Batch (class in taktiny.data) BatchMap (class in taktiny.data) BatchNorm (class in taktiny.nn) Bilinear (class in taktiny.nn) C Compose (class in taktiny.data) Conv (class in taktiny.nn) conv_general_dilated() (in module taktiny.utils.ops) ConvTranspose (class in taktiny.nn) D DataLoader (class in taktiny.data) DatasetConfig (class in taktiny.trainer) Dict (class in taktiny.nn) DoRAAdapter (class in taktiny.takt) DoRALinear (class in taktiny.nn) dot_general() (in module taktiny.utils.ops) Downsample (class in taktiny.nn) Dropout (class in taktiny.nn) E einsum() (in module taktiny.utils.ops) ELU (class in taktiny.nn) Embedding (class in taktiny.nn) eval() (taktiny.nn.Module method) evaluate() (taktiny.trainer.Trainer method) extract_params() (taktiny.trainer.Trainer method) F FeatureDropout (class in taktiny.nn) FeedForward (class in taktiny.nn) Filter (class in taktiny.data) flat_parameter_dict() (taktiny.nn.Module method) flat_state_dict() (taktiny.nn.Module method) FlatMap (class in taktiny.data) Flatten (class in taktiny.nn) Fold (class in taktiny.nn) format_bytes() (in module taktiny.utils.format) format_dtype() (in module taktiny.utils.format) format_params() (in module taktiny.utils.format) FractionalMaxPool (class in taktiny.nn) G GELU (class in taktiny.nn) get_logical_axis_rules() (in module taktiny.utils.spmd) GroupNorm (class in taktiny.nn) GRU (class in taktiny.nn) GRUCell (class in taktiny.nn) H HardSigmoid (class in taktiny.nn) HardSwish (class in taktiny.nn) HardTanh (class in taktiny.nn) I IndexMap (class in taktiny.data) initial_state() (taktiny.nn.GRUCell method) (taktiny.nn.LSTMCell method) (taktiny.nn.RNNCell method) items() (taktiny.nn.Dict method) K key (taktiny.nn.Rngs property) keys() (taktiny.nn.Dict method) L LayerNorm (class in taktiny.nn) Linear (class in taktiny.nn) linear() (in module taktiny.utils.ops) List (class in taktiny.nn) load_flat_state_dict() (taktiny.nn.Module method) load_state_dict() (taktiny.nn.Module method) logical_to_mesh_axes() (in module taktiny.utils.spmd) LoHaAdapter (class in taktiny.takt) LoHaLinear (class in taktiny.nn) LoKrAdapter (class in taktiny.takt) LoKrLinear (class in taktiny.nn) LoRAAdapter (class in taktiny.takt) LoRALinear (class in taktiny.nn) LPPool (class in taktiny.nn) LSTM (class in taktiny.nn) LSTMCell (class in taktiny.nn) M Map (class in taktiny.data) map_logical_axis_names() (in module taktiny.utils.spmd) MapFields (class in taktiny.data) mask_rank() (taktiny.nn.AdaLoRALinear method) MaxPool (class in taktiny.nn) MaxUnpool (class in taktiny.nn) Mish (class in taktiny.nn) Module (class in taktiny.nn) module() (in module taktiny.nn) O Optimizer (class in taktiny.takt.optimizer) orthogonal_loss() (taktiny.nn.AdaLoRALinear method) P Pack (class in taktiny.data) Padding (class in taktiny.nn) Parameter (class in taktiny.nn) parse_size() (in module taktiny.utils.format) prepare() (taktiny.takt.VeRAAdapter method) Pytree (class in taktiny.nn) Q quantization_rules() (in module taktiny.utils.quantization) quantize_conv_weight() (in module taktiny.utils.quantization) quantize_embedding_weight() (in module taktiny.utils.quantization) quantize_linear_weight() (in module taktiny.utils.quantization) quantize_model() (in module taktiny.utils.quantization) R RandomAccessSource (class in taktiny.data) RandomMap (class in taktiny.data) RecurrentBase (class in taktiny.nn) RecurrentLayer (class in taktiny.nn) ReLU (class in taktiny.nn) remove_callback() (taktiny.trainer.Trainer method) reset_running_stats() (taktiny.nn.BatchNorm method) RMSNorm (class in taktiny.nn) Rngs (class in taktiny.nn) RNN (class in taktiny.nn) RNNCell (class in taktiny.nn) S scan() (in module taktiny.utils.transforms) SELU (class in taktiny.nn) SeqStack (class in taktiny.nn) Sequential (class in taktiny.nn) set_logical_axis_rules() (in module taktiny.utils.spmd) Sigmoid (class in taktiny.nn) SiLU (class in taktiny.nn) SoftPlus (class in taktiny.nn) SoftSign (class in taktiny.nn) split_key() (taktiny.nn.Rngs method) split_rngs() (taktiny.nn.Rngs method) Stack (class in taktiny.nn) state_dict() (taktiny.nn.Module method) statistics() (taktiny.nn.BatchNorm method) StochasticDepth (class in taktiny.nn) Swish (class in taktiny.nn) T Takt (class in taktiny.takt) Tanh (class in taktiny.nn) TensorBoardCallback (class in taktiny.trainer) train() (taktiny.nn.Module method) (taktiny.trainer.Trainer method) train_validation_split() (in module taktiny.data) Trainer (class in taktiny.trainer) TrainerCallback (class in taktiny.trainer) TrainingConfig (class in taktiny.trainer) Transformer (class in taktiny.nn) TransformerDecoder (class in taktiny.nn) TransformerDecoderLayer (class in taktiny.nn) TransformerEncoder (class in taktiny.nn) TransformerEncoderLayer (class in taktiny.nn) U Unflatten (class in taktiny.nn) Unfold (class in taktiny.nn) update() (taktiny.takt.optimizer.Optimizer method) update_adapters() (taktiny.takt.Takt class method) update_running_stats() (taktiny.nn.BatchNorm method) Upsample (class in taktiny.nn) V value (taktiny.nn.Parameter property) values() (taktiny.nn.Dict method) VeRAAdapter (class in taktiny.takt) VeRALinear (class in taktiny.nn) vmap() (in module taktiny.utils.transforms) W WandbCallback (class in taktiny.trainer) with_logical_partitioning() (in module taktiny.utils.spmd)