API Reference

Binaries

• t5x.train binary
  • run_actions()
  • train()
• t5x.infer binary
  • FailFastThreadPoolExecutor
    • FailFastThreadPoolExecutor.check_for_exceptions()
    • FailFastThreadPoolExecutor.shutdown()
    • FailFastThreadPoolExecutor.submit()
  • SummarizeConfigFn
  • create_task_from_tfexample_file()
  • infer()
  • merge_chunks_to_file()
  • update_measurement_series()
  • write_inferences_to_file()
• t5x.eval binary
  • InferenceEvaluator
    • InferenceEvaluator.evaluate()
  • SummarizeConfigFn
  • evaluate()
• t5x.main binary
  • RunMode

Training

• t5x.trainer package
  • ActionMode
  • ArrayMapFuture
  • BaseAction
    • BaseAction.run()
  • BaseTrainer
    • BaseTrainer.close()
    • BaseTrainer.compile_eval()
    • BaseTrainer.compile_train()
    • BaseTrainer.eval()
    • BaseTrainer.train()
  • BaseTrainerConstructor
  • EarlyStoppingAction
    • EarlyStoppingAction.run()
  • LearningRateCallable
  • MetricValueMapFuture
  • MetricsManager
    • MetricsManager.start_duration_timer()
    • MetricsManager.summary_writer
    • MetricsManager.write_metrics_summary()
    • MetricsManager.write_scalar()
    • MetricsManager.write_scalars()
  • PartitionedEvalCallable
  • PartitionedTrainCallable
  • PreemptionError
  • SummarizeMetricsCallable
  • TerminateOnNanAction
    • TerminateOnNanAction.run()
  • TimeFuture
  • Trainer
  • WeightMetricsComputer
    • WeightMetricsComputer.compute_metrics()
  • accumulate_grads_microbatched()
  • apply_grads()
  • eval_step()
  • train_with_lr()
• t5x.optimizers package
  • MultiOptimizer
    • MultiOptimizer.apply_gradient()
    • MultiOptimizer.derive_logical_axes()
    • MultiOptimizer.set_param_axes()
    • MultiOptimizer.update_hyper_params()
  • OptaxStatePartitionRules
    • OptaxStatePartitionRules.derive_optax_logical_axes()
    • OptaxStatePartitionRules.derive_params_axes()
  • OptaxWrapper
    • OptaxWrapper.apply_gradient()
    • OptaxWrapper.derive_logical_axes()
    • OptaxWrapper.init_state()
    • OptaxWrapper.restore_state()
    • OptaxWrapper.state_dict()
  • Optimizer
    • Optimizer.optimizer_def
    • Optimizer.state
    • Optimizer.target
    • Optimizer.apply_gradient()
    • Optimizer.replace()
  • OptimizerDef
    • OptimizerDef.apply_gradient()
    • OptimizerDef.create()
    • OptimizerDef.restore_state()
    • OptimizerDef.update_hyper_params()
  • OptimizerDefType
  • OptimizerState
    • OptimizerState.replace()
  • OptimizerType
  • wrap_optax_optimizer()
• t5x.interactive_model package
  • InferenceType
  • InteractiveModel
    • InteractiveModel.evaluate()
    • InteractiveModel.evaluate_with_preprocessors()
    • InteractiveModel.infer_with_preprocessors()
    • InteractiveModel.predict_with_aux()
    • InteractiveModel.save_checkpoint()
    • InteractiveModel.score()
    • InteractiveModel.train_loop()
    • InteractiveModel.train_step()
    • InteractiveModel.train_step_from_batch_iterator()
    • InteractiveModel.train_step_with_preprocessors()
  • T5XScriptType
  • get_batches_from_seqio()
  • get_dataset_from_natural_text_examples()
  • get_gin_config_from_interactive_model()
  • get_seqio_task_from_examples()
• t5x.train_state package
  • FlaxOptimTrainState
    • FlaxOptimTrainState.replace()
  • InferenceState
    • InferenceState.param_states
    • InferenceState.replace()
  • TrainState
    • TrainState.apply_gradient()
    • TrainState.as_logical_axes()
    • TrainState.flax_mutables
    • TrainState.param_states
    • TrainState.params
    • TrainState.restore_state()
    • TrainState.state_dict()
    • TrainState.step
• t5x.state_utils package
  • apply_assignment_map()
  • flatten_state_dict()
  • get_name_tree()
  • intersect_state()
  • merge_state()
  • tensorstore_leaf()
• t5x.losses package
  • SpecialLossNormalizingFactor
    • SpecialLossNormalizingFactor.NUM_REAL_TARGET_TOKENS
    • SpecialLossNormalizingFactor.NUM_TOTAL_TARGET_TOKENS
    • SpecialLossNormalizingFactor.AVERAGE_PER_SEQUENCE
  • compute_weighted_cross_entropy()
  • convert_special_loss_normalizing_factor_to_enum()
  • get_loss_normalizing_factor_and_weights()
• t5x.metrics package
  • AveragePerStep
    • AveragePerStep.compute()
    • AveragePerStep.from_model_output()
    • AveragePerStep.merge()
    • AveragePerStep.replace()
  • Step
    • Step.compute()
    • Step.replace()
  • StepsPerTime
    • StepsPerTime.compute()
    • StepsPerTime.from_model_output()
    • StepsPerTime.merge()
    • StepsPerTime.replace()
  • Sum
    • Sum.compute()
    • Sum.from_model_output()
    • Sum.merge()
    • Sum.replace()
  • Time
    • Time.compute()
    • Time.merge()
    • Time.replace()
    • Time.replace_duration()
  • TimeRate
    • TimeRate.compute()
    • TimeRate.from_model_output()
    • TimeRate.merge()
    • TimeRate.replace()
  • create_metrics_dict()
  • set_step_metrics_num_steps()
  • set_time_metrics_duration()
  • shape_obj_to_defined_obj()
• t5x.utils package
  • CheckpointConfig
  • DatasetConfig
  • EvaluatorConstructor
  • GetDatasetCallable
  • GetEvalDatasetCallable
  • InferFnCallable
  • InferStepWithRngCallable
  • InferStepWithoutRngCallable
  • InitFnCallable
  • LearningRateCallable
  • LegacyCheckpointManager
    • LegacyCheckpointManager.restore()
    • LegacyCheckpointManager.save()
  • LegacyCheckpointer
    • LegacyCheckpointer.restore()
    • LegacyCheckpointer.save()
  • RestoreCheckpointConfig
    • RestoreCheckpointConfig.checkpoint_manager_cls
    • RestoreCheckpointConfig.checkpointer_cls
  • SaveCheckpointConfig
    • SaveCheckpointConfig.checkpoint_manager_cls
    • SaveCheckpointConfig.checkpointer_cls
  • ShardedDatasetIterator
    • ShardedDatasetIterator.element_spec
    • ShardedDatasetIterator.reset()
    • ShardedDatasetIterator.restore()
    • ShardedDatasetIterator.save()
  • TrainStateInitializer
    • TrainStateInitializer.global_train_state_shape
    • TrainStateInitializer.train_state_axes
    • TrainStateInitializer.from_checkpoint()
    • TrainStateInitializer.from_checkpoint_or_scratch()
    • TrainStateInitializer.from_checkpoints()
    • TrainStateInitializer.from_scratch()
  • create_checkpoint_manager_and_restore()
  • create_learning_rate_scheduler()
  • create_orbax_checkpoint_manager()
  • find_first_checkpoint_step()
  • find_next_checkpoint_step()
  • flatten_dict_string_keys()
  • get_dataset()
  • get_dataset_inner()
  • get_fallback_state()
  • get_first_valid_restore_config_and_paths()
  • get_infer_fn()
  • get_training_eval_datasets()
  • get_vocabulary()
  • get_zeros_batch_like_dataset()
  • import_module()
  • log_model_info()
  • multihost_assert_equal()
  • override_params_axes_names()
  • prepare_train_iter()
  • restore()
  • round_vocab_size_to_multiple()
  • sync_global_devices()
  • verify_matching_vocabs()
• t5x.adafactor package
  • Adafactor
    • Adafactor.apply_gradient()
    • Adafactor.derive_logical_axes()
    • Adafactor.set_param_axes()
  • FactorDim
  • HParamMap
  • HeuristicRule

Inference

• t5x.decoding package
  • BeamState
    • BeamState.replace()
  • DecodingState
    • DecodingState.cur_index
    • DecodingState.sequences
    • DecodingState.cur_token
    • DecodingState.cache
    • DecodingState.replace()
  • SamplingLoopState
    • SamplingLoopState.step
    • SamplingLoopState.cur_index
    • SamplingLoopState.sequences
    • SamplingLoopState.cache
    • SamplingLoopState.cur_token
    • SamplingLoopState.ended
    • SamplingLoopState.rng
    • SamplingLoopState.log_prob
    • SamplingLoopState.replace()
  • add_beam_dim()
  • beam_init()
  • beam_search()
  • brevity_penalty()
  • cache_gather_beams()
  • cache_map()
  • flat_batch_beam_expand()
  • flatten_beam_dim()
  • gather_beams()
  • gather_topk_beams()
  • temperature_sample()
  • top_k_two_stage()
  • unflatten_beam_dim()

Models

• t5x.models package
  • BaseModel
    • BaseModel.eval_fn()
    • BaseModel.get_initial_variables()
    • BaseModel.loss_fn()
    • BaseModel.predict_batch()
    • BaseModel.predict_batch_with_aux()
    • BaseModel.score_batch()
  • BaseTransformerModel
    • BaseTransformerModel.loss_fn()
  • DecodeFnCallable
  • DecoderOnlyModel
    • DecoderOnlyModel.FEATURE_CONVERTER_CLS
    • DecoderOnlyModel.get_initial_variables()
    • DecoderOnlyModel.predict_batch_with_aux()
    • DecoderOnlyModel.score_batch()
  • DecoderParams
  • EncoderDecoderModel
    • EncoderDecoderModel.FEATURE_CONVERTER_CLS
    • EncoderDecoderModel.get_initial_variables()
    • EncoderDecoderModel.predict_batch_with_aux()
    • EncoderDecoderModel.score_batch()
  • TokensIdsToLogitsCallable
  • compute_base_metrics()
  • compute_metrics()
  • compute_weighted_accuracy()
  • count_packed_examples()
  • remove_prefix()

Checkpointing

• t5x.checkpoints package
  • CheckpointManagerConstructor
  • Checkpointer
    • Checkpointer.checkpoints_dir
    • Checkpointer.keep
    • Checkpointer.restore_dtype
    • Checkpointer.save_dtype
    • Checkpointer.keep_dataset_checkpoints
    • Checkpointer.all_dataset_checkpoint_steps()
    • Checkpointer.all_steps()
    • Checkpointer.convert_from_tf_checkpoint()
    • Checkpointer.latest_step()
    • Checkpointer.restore()
    • Checkpointer.restore_from_tf_checkpoint()
    • Checkpointer.save()
  • CheckpointerConstructor
  • DatasetCheckpointHandler
    • DatasetCheckpointHandler.restore()
    • DatasetCheckpointHandler.save()
    • DatasetCheckpointHandler.structure()
  • OrbaxCheckpointManagerInterface
    • OrbaxCheckpointManagerInterface.restore()
    • OrbaxCheckpointManagerInterface.restore_from_tf_checkpoint()
    • OrbaxCheckpointManagerInterface.save()
  • RestoreStateTransformationFn
  • SaveBestCheckpointer
    • SaveBestCheckpointer.checkpoints_dir
    • SaveBestCheckpointer.keep
    • SaveBestCheckpointer.restore_dtype
    • SaveBestCheckpointer.save_dtype
    • SaveBestCheckpointer.metric_name_to_monitor
    • SaveBestCheckpointer.metric_mode
    • SaveBestCheckpointer.keep_checkpoints_without_metrics
    • SaveBestCheckpointer.force_keep_period
    • SaveBestCheckpointer.keep_dataset_checkpoints
  • SaveStateTransformationFn
  • all_dataset_checkpoint_steps()
  • all_steps()
  • fake_param_info()
  • find_checkpoint()
  • get_checkpoint_dir()
  • get_local_data()
  • get_step_from_checkpoint_dir()
  • latest_step()
  • load_t5x_checkpoint()
  • populate_metrics_for_steps()
• t5x.checkpoint_utils package
  • CheckpointTypes
  • detect_checkpoint_type()
  • is_pinned_checkpoint()
  • pin_checkpoint()
  • pinned_checkpoint_filepath()
  • remove_checkpoint_dir()
  • remove_dataset_checkpoint()
  • unpin_checkpoint()
• t5x.checkpoint_importer package
  • CheckpointTranslator
    • CheckpointTranslator.add()
    • CheckpointTranslator.apply()
  • LazyArray
  • LazyAwaitableArray
    • LazyAwaitableArray.from_array()
    • LazyAwaitableArray.from_tensor_store_spec()
    • LazyAwaitableArray.from_tensor_store_spec_or_array()
  • LazyThreadPoolArray
  • attention_layers()
  • final_layernorms()
  • layernorms()
  • load_tf_ckpt()
  • mlpblock()
  • rel_embeddings()
  • restore_from_t5_checkpoint()

Paritioning

• t5x.partitioning package
  • AxisNames
  • BasePartitioner
    • BasePartitioner.compile()
    • BasePartitioner.data_mesh_size
    • BasePartitioner.data_shard_id
    • BasePartitioner.data_shards
    • BasePartitioner.get_data_layout()
    • BasePartitioner.get_local_chunk_info()
    • BasePartitioner.get_logical_axes()
    • BasePartitioner.get_mesh_axes()
    • BasePartitioner.move_params_to_devices()
    • BasePartitioner.partition()
  • BasePjitPartitioner
    • BasePjitPartitioner.compile()
    • BasePjitPartitioner.partition()
  • DataLayout
  • LocalChunkInfo
  • LocalChunker
    • LocalChunker.get_local_chunk_info()
    • LocalChunker.get_replica_id()
    • LocalChunker.get_shard_id()
  • PjitPartitioner
    • PjitPartitioner.get_logical_axes()
    • PjitPartitioner.get_mesh_axes()
    • PjitPartitioner.logical_axis_rules
    • PjitPartitioner.partition()
  • PjittedFnWithContext
  • bounds_from_last_device()
  • default_mesh()
  • get_coords()
  • get_cpu_mesh()
  • get_gpu_mesh()
  • get_mesh()
  • global_mesh_defined()
  • standard_logical_axis_rules()
  • with_sharding_constraint()

Config

• t5x.config_utils package
  • config_with_fiddle()
  • direct_summarize_fiddle_config()
  • prepare_to_summarize()
  • run()
  • sanitize_summary_getter()
  • using_fdl()
• t5x.gin_utils package
  • bool_fn()
  • get_gin_config_str()
  • parse_gin_flags()
  • rewrite_gin_args()
  • run()
  • string_split_fn()
  • sum_fn()
  • summarize_gin_config()

Utils

• t5x.test_utils package
• t5x.binary_search package
  • float32_bsearch()
  • int32_bsearch()
  • topk_mask()
  • topp_mask()