Index A | B | C | D | E | F | G | H | I | K | L | M | N | O | P | Q | R | S | T | U | V | W A activation() (flax_extra.layer.FeedForward method) add() (in module flax_extra.combinator) aggregation (flax_extra.layer.Encoding attribute) apply (flax_extra.evaluation.EvalTask attribute) (flax_extra.training.TrainTask attribute) attend() (in module flax_extra.layer) Attention (class in flax_extra.layer) attention (flax_extra.model.perceiver.CrossAttentionBlock attribute) (flax_extra.model.perceiver.Decoder attribute) (flax_extra.model.perceiver.Encoder attribute) (flax_extra.model.perceiver.Processor attribute) (flax_extra.model.perceiver.SelfAttentionBlock attribute) axis (flax_extra.combinator.Concatenate attribute) (flax_extra.operator.ShiftRight attribute) B batch_per_device() (in module flax_extra.util) bias_init() (flax_extra.layer.Attention method) bindings (flax_extra.operator.Rearrange attribute) branch() (in module flax_extra.combinator) bytes_tokenizer() (in module flax_extra.data) BytesTokenizer (class in flax_extra.data) C Checkpoint (class in flax_extra.checkpoint) CheckpointFile (class in flax_extra.checkpoint) CheckpointFileReader (class in flax_extra.checkpoint) CheckpointFileWriter (class in flax_extra.checkpoint) children (flax_extra.combinator.Parallel attribute) (flax_extra.combinator.Serial attribute) children_signatures (flax_extra.combinator.Parallel attribute) (flax_extra.combinator.Serial attribute) Combinator (class in flax_extra.combinator) Concatenate (class in flax_extra.combinator) concatenate() (in module flax_extra.combinator) cross_attention_mask() (in module flax_extra.layer) CrossAttentionBlock (class in flax_extra.model.perceiver) D d_hidden (flax_extra.layer.LSTM attribute) d_max (flax_extra.layer.TrainablePositionalPadding attribute) d_model (flax_extra.model.RNNLM attribute) d_output() (flax_extra.layer.Attention method) d_qk() (flax_extra.layer.Attention method) d_reserved (flax_extra.data.BytesTokenizer property) (flax_extra.layer.MultimodalEncoding attribute) (flax_extra.layer.MultimodalPositionalEncoding attribute) d_v() (flax_extra.layer.Attention method) data (flax_extra.evaluation.EvalTask attribute) (flax_extra.training.TrainTask attribute) Decoder (class in flax_extra.model.perceiver) decoder_attention (flax_extra.model.PerceiverIO attribute) decoder_feed_forward (flax_extra.model.PerceiverIO attribute) decoder_query_encoding (flax_extra.model.PerceiverIO attribute) Decoding (class in flax_extra.layer) default_metric_value() (flax_extra.checkpoint.HighestCheckpointFileWriter method) (flax_extra.checkpoint.LowestCheckpointFileWriter method) deterministic (flax_extra.layer.Attention attribute) (flax_extra.model.perceiver.CrossAttentionBlock attribute) (flax_extra.model.perceiver.Decoder attribute) (flax_extra.model.perceiver.Encoder attribute) (flax_extra.model.perceiver.Processor attribute) (flax_extra.model.perceiver.SelfAttentionBlock attribute) (flax_extra.model.PerceiverIO attribute) (flax_extra.model.RNNLM attribute) dimension (flax_extra.layer.TrainablePositionalEncoding attribute) dir (flax_extra.checkpoint.CheckpointFileReader attribute) div() (in module flax_extra.combinator) Drop (class in flax_extra.combinator) drop() (in module flax_extra.combinator) dropout_rate (flax_extra.layer.Attention attribute) (flax_extra.model.perceiver.CrossAttentionBlock attribute) (flax_extra.model.perceiver.Decoder attribute) (flax_extra.model.perceiver.Encoder attribute) (flax_extra.model.perceiver.Processor attribute) (flax_extra.model.perceiver.SelfAttentionBlock attribute) (flax_extra.model.RNNLM attribute) dtype (flax_extra.layer.TrainablePositionalEncoding attribute) Dup (class in flax_extra.combinator) dup() (in module flax_extra.combinator) E elapsed_time (flax_extra.checkpoint.Checkpoint attribute) EmbedDecoding (class in flax_extra.layer) embedding (flax_extra.layer.EmbedDecoding attribute) embedding_decoding() (flax_extra.layer.Decoding method) Encoder (class in flax_extra.model.perceiver) encoder_attention (flax_extra.model.PerceiverIO attribute) encoder_feed_forward (flax_extra.model.PerceiverIO attribute) encoder_query_encoding (flax_extra.model.PerceiverIO attribute) Encoding (class in flax_extra.layer) EvalLoop (class in flax_extra.evaluation) EvalTask (class in flax_extra.evaluation) F feed_forward (flax_extra.model.perceiver.CrossAttentionBlock attribute) (flax_extra.model.perceiver.Decoder attribute) (flax_extra.model.perceiver.Encoder attribute) (flax_extra.model.perceiver.Processor attribute) (flax_extra.model.perceiver.SelfAttentionBlock attribute) FeedForward (class in flax_extra.layer) flax_extra module flax_extra.batch module flax_extra.checkpoint module flax_extra.combinator module flax_extra.console module flax_extra.data module flax_extra.evaluation module flax_extra.layer module flax_extra.layer.io module flax_extra.model module flax_extra.model.perceiver module flax_extra.model.rnn module flax_extra.operator module flax_extra.random module flax_extra.training module flax_extra.util module flax_extra.version module fold() (in module flax_extra.combinator) Foldl (class in flax_extra.combinator) foldl() (in module flax_extra.combinator) FourierPositionEncoding (class in flax_extra.layer) func (flax_extra.combinator.Foldl attribute) G grads (flax_extra.checkpoint.Checkpoint attribute) group (flax_extra.checkpoint.HighestCheckpointFileReader attribute) (flax_extra.checkpoint.LowestCheckpointFileReader attribute) H HighestCheckpointFileReader (class in flax_extra.checkpoint) HighestCheckpointFileWriter (class in flax_extra.checkpoint) I Identity (class in flax_extra.combinator) identity() (in module flax_extra.combinator) indices (flax_extra.combinator.Select attribute) initial_state() (flax_extra.layer.LSTM method) input_encoding (flax_extra.model.PerceiverIO attribute) input_encoding() (in module flax_extra.layer.io) into_collection() (in module flax_extra.random) into_sequence() (in module flax_extra.random) is_best() (flax_extra.checkpoint.HighestCheckpointFileWriter method) (flax_extra.checkpoint.LowestCheckpointFileWriter method) K keep (flax_extra.checkpoint.CheckpointFileWriter attribute) kernel_init() (flax_extra.layer.Attention method) (flax_extra.layer.FeedForward method) (flax_extra.layer.TrainablePositionalEncoding method) L log() (in module flax_extra.console) log_step() (in module flax_extra.console) loss (flax_extra.checkpoint.Checkpoint attribute) (flax_extra.training.TrainTask attribute) LowestCheckpointFileReader (class in flax_extra.checkpoint) LowestCheckpointFileWriter (class in flax_extra.checkpoint) LSTM (class in flax_extra.layer) LSTMCell (class in flax_extra.layer) M mask_rates (flax_extra.layer.MultimodalEncoding attribute) metric (flax_extra.checkpoint.HighestCheckpointFileReader attribute) (flax_extra.checkpoint.LowestCheckpointFileReader attribute) metrics (flax_extra.checkpoint.CheckpointFile attribute) (flax_extra.checkpoint.Summary attribute) (flax_extra.evaluation.EvalTask attribute) modalities (flax_extra.layer.MultimodalDecoding attribute) (flax_extra.layer.MultimodalEncoding attribute) (flax_extra.layer.MultimodalPositionalEncoding attribute) model_params (flax_extra.checkpoint.Checkpoint attribute) (flax_extra.checkpoint.CheckpointFile attribute) model_state (flax_extra.checkpoint.Checkpoint attribute) (flax_extra.checkpoint.CheckpointFile attribute) module flax_extra flax_extra.batch flax_extra.checkpoint flax_extra.combinator flax_extra.console flax_extra.data flax_extra.evaluation flax_extra.layer flax_extra.layer.io flax_extra.model flax_extra.model.perceiver flax_extra.model.rnn flax_extra.operator flax_extra.random flax_extra.training flax_extra.util flax_extra.version mul() (in module flax_extra.combinator) multimodal_embedding_decoding() (flax_extra.layer.MultimodalDecoding method) MultimodalDecoding (class in flax_extra.layer) MultimodalEncoding (class in flax_extra.layer) MultimodalPositionalEncoding (class in flax_extra.layer) N n_bands (flax_extra.layer.FourierPositionEncoding attribute) n_blocks (flax_extra.model.perceiver.Processor attribute) n_completed_steps (flax_extra.checkpoint.Checkpoint attribute) n_heads (flax_extra.layer.Attention attribute) n_layers (flax_extra.model.RNNLM attribute) n_modalities (flax_extra.layer.MultimodalDecoding property) (flax_extra.layer.MultimodalEncoding property) (flax_extra.layer.MultimodalPositionalEncoding property) n_positions (flax_extra.operator.ShiftRight attribute) n_processor_blocks (flax_extra.model.PerceiverIO attribute) n_processor_shards (flax_extra.model.PerceiverIO attribute) n_shards (flax_extra.model.perceiver.Processor attribute) n_steps (flax_extra.evaluation.EvalLoop property) (flax_extra.training.TrainLoop property) n_steps_per_checkpoint (flax_extra.training.TrainLoop property) name (flax_extra.layer.Attention attribute) (flax_extra.layer.Decoding attribute) (flax_extra.layer.EmbedDecoding attribute) (flax_extra.layer.Encoding attribute) (flax_extra.layer.FeedForward attribute) (flax_extra.layer.FourierPositionEncoding attribute) (flax_extra.layer.LSTM attribute) (flax_extra.layer.LSTMCell attribute) (flax_extra.layer.MultimodalDecoding attribute) (flax_extra.layer.MultimodalEncoding attribute) (flax_extra.layer.MultimodalPositionalEncoding attribute) (flax_extra.layer.SelfAttention attribute) (flax_extra.layer.TrainablePositionalEncoding attribute) (flax_extra.layer.TrainablePositionalMasking attribute) (flax_extra.layer.TrainablePositionalPadding attribute) (flax_extra.model.perceiver.CrossAttentionBlock attribute) (flax_extra.model.perceiver.Decoder attribute) (flax_extra.model.perceiver.Encoder attribute) (flax_extra.model.perceiver.Processor attribute) (flax_extra.model.perceiver.SelfAttentionBlock attribute) (flax_extra.model.PerceiverIO attribute) (flax_extra.model.RNNLM attribute) next_checkpoint() (flax_extra.training.TrainLoop method) next_step() (flax_extra.training.TrainLoop method) normalize_batch() (in module flax_extra.batch) normalize_batch_per_device() (in module flax_extra.batch) O optimizer (flax_extra.training.TrainTask attribute) optimizer_state (flax_extra.checkpoint.Checkpoint attribute) (flax_extra.checkpoint.CheckpointFile attribute) originate() (in module flax_extra.util) output_decoding (flax_extra.model.PerceiverIO attribute) output_decoding() (in module flax_extra.layer.io) output_dir (flax_extra.checkpoint.CheckpointFileWriter attribute) (flax_extra.checkpoint.HighestCheckpointFileWriter attribute) (flax_extra.checkpoint.LowestCheckpointFileWriter attribute) (flax_extra.checkpoint.SummaryWriter attribute) overwrite (flax_extra.checkpoint.CheckpointFileWriter attribute) P pad() (flax_extra.data.BytesTokenizer method) pad_id (flax_extra.operator.ShiftRight attribute) Parallel (class in flax_extra.combinator) parallel() (in module flax_extra.combinator) parent (flax_extra.layer.Attention attribute) (flax_extra.layer.Decoding attribute) (flax_extra.layer.EmbedDecoding attribute) (flax_extra.layer.Encoding attribute) (flax_extra.layer.FeedForward attribute) (flax_extra.layer.FourierPositionEncoding attribute) (flax_extra.layer.LSTM attribute) (flax_extra.layer.LSTMCell attribute) (flax_extra.layer.MultimodalDecoding attribute) (flax_extra.layer.MultimodalEncoding attribute) (flax_extra.layer.MultimodalPositionalEncoding attribute) (flax_extra.layer.SelfAttention attribute) (flax_extra.layer.TrainablePositionalEncoding attribute) (flax_extra.layer.TrainablePositionalMasking attribute) (flax_extra.layer.TrainablePositionalPadding attribute) (flax_extra.model.perceiver.CrossAttentionBlock attribute) (flax_extra.model.perceiver.Decoder attribute) (flax_extra.model.perceiver.Encoder attribute) (flax_extra.model.perceiver.Processor attribute) (flax_extra.model.perceiver.SelfAttentionBlock attribute) (flax_extra.model.PerceiverIO attribute) (flax_extra.model.RNNLM attribute) pattern (flax_extra.operator.Rearrange attribute) PerceiverIO (class in flax_extra.model) positional_encoding (flax_extra.layer.Encoding attribute) postprocessing() (flax_extra.layer.Decoding method) precision (flax_extra.layer.Attention attribute) (flax_extra.model.perceiver.CrossAttentionBlock attribute) (flax_extra.model.perceiver.Decoder attribute) (flax_extra.model.perceiver.Encoder attribute) (flax_extra.model.perceiver.Processor attribute) (flax_extra.model.perceiver.SelfAttentionBlock attribute) (flax_extra.model.PerceiverIO attribute) preprocessing() (flax_extra.layer.Encoding method) Processor (class in flax_extra.model.perceiver) processor_attention (flax_extra.model.PerceiverIO attribute) processor_feed_forward (flax_extra.model.PerceiverIO attribute) Q QKVAttention (in module flax_extra.layer) query_encoding() (in module flax_extra.layer.io) R rate (flax_extra.layer.TrainablePositionalMasking attribute) read() (flax_extra.checkpoint.CheckpointFileReader method) Rearrange (class in flax_extra.operator) replace() (flax_extra.checkpoint.CheckpointFile method) reserved_ids (flax_extra.data.BytesTokenizer attribute) ReshapeBatch (class in flax_extra.operator) residual() (in module flax_extra.combinator) rnn_type (flax_extra.model.RNNLM attribute) RNNLM (class in flax_extra.model) run() (flax_extra.training.TrainLoop method) S scope (flax_extra.layer.Attention attribute) (flax_extra.layer.Decoding attribute) (flax_extra.layer.EmbedDecoding attribute) (flax_extra.layer.Encoding attribute) (flax_extra.layer.FeedForward attribute) (flax_extra.layer.FourierPositionEncoding attribute) (flax_extra.layer.LSTM attribute) (flax_extra.layer.LSTMCell attribute) (flax_extra.layer.MultimodalDecoding attribute) (flax_extra.layer.MultimodalEncoding attribute) (flax_extra.layer.MultimodalPositionalEncoding attribute) (flax_extra.layer.SelfAttention attribute) (flax_extra.layer.TrainablePositionalEncoding attribute) (flax_extra.layer.TrainablePositionalMasking attribute) (flax_extra.layer.TrainablePositionalPadding attribute) (flax_extra.model.perceiver.CrossAttentionBlock attribute) (flax_extra.model.perceiver.Decoder attribute) (flax_extra.model.perceiver.Encoder attribute) (flax_extra.model.perceiver.Processor attribute) (flax_extra.model.perceiver.SelfAttentionBlock attribute) (flax_extra.model.PerceiverIO attribute) (flax_extra.model.RNNLM attribute) Select (class in flax_extra.combinator) select() (in module flax_extra.combinator) SelfAttention (class in flax_extra.layer) SelfAttentionBlock (class in flax_extra.model.perceiver) seqlen (flax_extra.layer.TrainablePositionalEncoding attribute) seqshape (flax_extra.layer.FourierPositionEncoding attribute) sequence() (in module flax_extra.random) Serial (class in flax_extra.combinator) serial() (in module flax_extra.combinator) shape (flax_extra.operator.ReshapeBatch attribute) ShiftRight (class in flax_extra.operator) signature (flax_extra.combinator.Combinator attribute) (flax_extra.combinator.Drop attribute) (flax_extra.combinator.Dup attribute) (flax_extra.combinator.Identity attribute) spatial_coordinate_bounds (flax_extra.layer.FourierPositionEncoding attribute) split_per_device() (in module flax_extra.random) step (flax_extra.checkpoint.Checkpoint attribute) (flax_extra.checkpoint.CheckpointFile attribute) (flax_extra.training.TrainLoop property) sub() (in module flax_extra.combinator) Summary (class in flax_extra.checkpoint) SummaryLogger (class in flax_extra.checkpoint) SummaryWriter (class in flax_extra.checkpoint) T target (flax_extra.checkpoint.CheckpointFileReader attribute) target_encoding() (in module flax_extra.layer.io) to_ids() (flax_extra.data.BytesTokenizer method) to_tokens() (flax_extra.data.BytesTokenizer method) TrainablePositionalEncoding (class in flax_extra.layer) TrainablePositionalMasking (class in flax_extra.layer) TrainablePositionalPadding (class in flax_extra.layer) TrainLoop (class in flax_extra.training) TrainTask (class in flax_extra.training) U use_bias (flax_extra.layer.Attention attribute) use_decoder_q_residual (flax_extra.model.PerceiverIO attribute) use_encoder_q_residual (flax_extra.model.PerceiverIO attribute) use_q_residual (flax_extra.model.perceiver.CrossAttentionBlock attribute) (flax_extra.model.perceiver.Decoder attribute) (flax_extra.model.perceiver.Encoder attribute) use_sine_only (flax_extra.layer.FourierPositionEncoding attribute) use_spatial_coordinates (flax_extra.layer.FourierPositionEncoding attribute) V version (flax_extra.checkpoint.CheckpointFile attribute) vocab_size (flax_extra.data.BytesTokenizer property) (flax_extra.model.RNNLM attribute) W widening_factor (flax_extra.layer.FeedForward attribute) write() (flax_extra.checkpoint.CheckpointFileWriter method)