├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── RELEASE.md ├── WORKSPACE ├── praxis ├── AUTHORS ├── BUILD ├── __init__.py ├── asserts.py ├── asserts_test.py ├── base_hyperparams.py ├── base_hyperparams_test.py ├── base_input.py ├── base_input_test.py ├── base_layer.py ├── base_layer_test.py ├── base_model.py ├── beam_search.py ├── beam_search_test.py ├── build-visibility.bzl ├── contrib │ └── gpu │ │ └── scripts_gpu │ │ ├── lora_layers.py │ │ └── models.py ├── decoder_hparams.py ├── decoder_utils.py ├── decoder_utils_test.py ├── fiddle_tags.py ├── flat_beam_search.py ├── flat_beam_search_test.py ├── flax_utils.py ├── gshard_utils.py ├── layers │ ├── BUILD │ ├── __init__.py │ ├── activations.py │ ├── activations_test.py │ ├── adapters.py │ ├── adapters_test.py │ ├── attentions.py │ ├── attentions_test.py │ ├── augmentations.py │ ├── augmentations_test.py │ ├── base_ops.py │ ├── bregman.py │ ├── bregman_test.py │ ├── chain │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── chain.py │ │ ├── chain_extensions.py │ │ └── chain_test.py │ ├── checkpoint_policy.py │ ├── chunk.py │ ├── chunk_test.py │ ├── conformers.py │ ├── conformers_test.py │ ├── convolutions.py │ ├── convolutions_test.py │ ├── ctc_objectives.py │ ├── ctc_objectives_test.py │ ├── einsum.py │ ├── einsum_test.py │ ├── embedding_softmax.py │ ├── embedding_softmax_test.py │ ├── flax_adapter.py │ ├── flax_adapter_test.py │ ├── frnn.py │ ├── frnn_test.py │ ├── glam.py │ ├── gpu_fast_attention.py │ ├── grok.py │ ├── grouped_query_attention.py │ ├── grouped_query_attention_test.py │ ├── injection │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── fp8_nvidia_gpu.py │ │ └── fp8_nvidia_gpu_test.py │ ├── linears.py │ ├── linears_test.py │ ├── losses.py │ ├── mobilenet.py │ ├── models.py │ ├── models_test.py │ ├── multi_query_attention.py │ ├── multi_query_attention_test.py │ ├── ngrammer.py │ ├── ngrammer_test.py │ ├── normalizations.py │ ├── normalizations_test.py │ ├── pipeline.py │ ├── pipeline_test.py │ ├── poolings.py │ ├── poolings_test.py │ ├── quantization │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── attentions.py │ │ ├── attentions_test.py │ │ ├── automl_select.py │ │ ├── automl_select_test.py │ │ ├── conformers.py │ │ ├── conformers_test.py │ │ ├── convolutions.py │ │ ├── convolutions_test.py │ │ ├── einsum.py │ │ ├── einsum_test.py │ │ ├── embedding_softmax.py │ │ ├── embedding_softmax_test.py │ │ ├── linears.py │ │ ├── linears_test.py │ │ ├── multi_query_attention.py │ │ ├── multi_query_attention_test.py │ │ ├── ngrammer.py │ │ ├── ngrammer_test.py │ │ ├── operations.py │ │ ├── operations_test.py │ │ ├── optimization.py │ │ ├── optimization_test.py │ │ ├── overflow_check.py │ │ ├── overflow_check_test.py │ │ ├── quantization_hparams.py │ │ ├── quantization_test.py │ │ ├── quantize.py │ │ ├── quantize_test.py │ │ ├── quantizer.py │ │ ├── quantizer_test.py │ │ ├── searchable.py │ │ ├── searchable_test.py │ │ ├── sparsity │ │ │ ├── BUILD │ │ │ ├── __init__.py │ │ │ ├── attentions_test.py │ │ │ ├── linears_test.py │ │ │ ├── sparsifier.py │ │ │ ├── sparsifier_test.py │ │ │ ├── sparsity.py │ │ │ ├── sparsity_hparams.py │ │ │ ├── sparsity_modes.py │ │ │ └── sparsity_test.py │ │ ├── tests │ │ │ ├── attention_projection_aqt_test.py │ │ │ ├── attention_projection_fq_test.py │ │ │ ├── attention_projection_ptq_test.py │ │ │ ├── combined_qkv_projection_aqt_test.py │ │ │ ├── combined_qkv_projection_fq_test.py │ │ │ ├── combined_qkv_projection_ptq_test.py │ │ │ ├── dotproduct_attention_aqt_test.py │ │ │ ├── linear_ptq_test.py │ │ │ ├── one_headed_attention_projection_aqt_test.py │ │ │ └── test_util.py │ │ ├── utils.py │ │ └── utils_test.py │ ├── quantizer.py │ ├── quantizer_objectives.py │ ├── quantizer_objectives_test.py │ ├── quantizer_test.py │ ├── repeats.py │ ├── repeats_test.py │ ├── resnets.py │ ├── rnn_cell.py │ ├── rnn_cell_test.py │ ├── searchable.py │ ├── searchable_test.py │ ├── sedd │ │ └── sedd.ipynb │ ├── sequential.py │ ├── sequential_test.py │ ├── sharding.py │ ├── shared_layers_test.py │ ├── spectrum_augmenter.py │ ├── spectrum_augmenter_test.py │ ├── ssm.py │ ├── ssm_test.py │ ├── ssm_transformers.py │ ├── ssm_transformers_test.py │ ├── stats.py │ ├── stats_test.py │ ├── stochastics.py │ ├── stochastics_test.py │ ├── test_layers.py │ ├── transformer_models.py │ ├── transformer_models_encoder_decoder_test.py │ ├── transformer_models_test.py │ ├── transformers.py │ ├── transformers_test.py │ ├── vanillanets.py │ ├── vanillanets_test.py │ ├── video │ │ ├── BUILD │ │ ├── __init__.py │ │ ├── enc_dec_3dcnn.py │ │ ├── enc_dec_3dcnn_test.py │ │ ├── losses.py │ │ ├── losses_test.py │ │ ├── quantizer.py │ │ ├── quantizer_test.py │ │ ├── vqvae.py │ │ └── vqvae_test.py │ ├── vits.py │ └── vits_test.py ├── lazy_loader.py ├── lingvo_lib.py ├── metric_utils.py ├── optimizer_prefix_vectorization.py ├── optimizer_prefix_vectorization_test.py ├── optimizers.py ├── optimizers_test.py ├── pax_fiddle.py ├── pax_fiddle_test.py ├── pip_package │ ├── build_pip_pkg.sh │ ├── cloudbuild-postsubmit.yaml │ ├── cloudbuild-presubmit.yaml │ ├── cloudbuild-release.yaml │ ├── postsubmit.Dockerfile │ ├── prepare_release.sh │ ├── presubmit.Dockerfile │ ├── release.Dockerfile │ └── requirements.txt ├── praxis.bzl ├── py_utils.py ├── py_utils_test.py ├── pytypes.py ├── pytypes_test.py ├── sample_decode.py ├── sample_decode_test.py ├── schedules.py ├── schedules_test.py ├── test_utils.py ├── token_samplers.py ├── token_samplers_test.py ├── trees.py └── trees_test.py ├── requirements.in └── setup.py /CONTRIBUTING.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/CONTRIBUTING.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/README.md -------------------------------------------------------------------------------- /RELEASE.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/RELEASE.md -------------------------------------------------------------------------------- /WORKSPACE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/WORKSPACE -------------------------------------------------------------------------------- /praxis/AUTHORS: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/AUTHORS -------------------------------------------------------------------------------- /praxis/BUILD: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/BUILD -------------------------------------------------------------------------------- /praxis/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/__init__.py -------------------------------------------------------------------------------- /praxis/asserts.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/asserts.py -------------------------------------------------------------------------------- /praxis/asserts_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/asserts_test.py -------------------------------------------------------------------------------- /praxis/base_hyperparams.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/base_hyperparams.py -------------------------------------------------------------------------------- /praxis/base_hyperparams_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/base_hyperparams_test.py -------------------------------------------------------------------------------- /praxis/base_input.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/base_input.py -------------------------------------------------------------------------------- /praxis/base_input_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/base_input_test.py -------------------------------------------------------------------------------- /praxis/base_layer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/base_layer.py -------------------------------------------------------------------------------- /praxis/base_layer_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/base_layer_test.py -------------------------------------------------------------------------------- /praxis/base_model.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/base_model.py -------------------------------------------------------------------------------- /praxis/beam_search.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/beam_search.py -------------------------------------------------------------------------------- /praxis/beam_search_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/beam_search_test.py -------------------------------------------------------------------------------- /praxis/build-visibility.bzl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/build-visibility.bzl -------------------------------------------------------------------------------- /praxis/contrib/gpu/scripts_gpu/lora_layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/contrib/gpu/scripts_gpu/lora_layers.py -------------------------------------------------------------------------------- /praxis/contrib/gpu/scripts_gpu/models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/contrib/gpu/scripts_gpu/models.py -------------------------------------------------------------------------------- /praxis/decoder_hparams.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/decoder_hparams.py -------------------------------------------------------------------------------- /praxis/decoder_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/decoder_utils.py -------------------------------------------------------------------------------- /praxis/decoder_utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/decoder_utils_test.py -------------------------------------------------------------------------------- /praxis/fiddle_tags.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/fiddle_tags.py -------------------------------------------------------------------------------- /praxis/flat_beam_search.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/flat_beam_search.py -------------------------------------------------------------------------------- /praxis/flat_beam_search_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/flat_beam_search_test.py -------------------------------------------------------------------------------- /praxis/flax_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/flax_utils.py -------------------------------------------------------------------------------- /praxis/gshard_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/gshard_utils.py -------------------------------------------------------------------------------- /praxis/layers/BUILD: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/BUILD -------------------------------------------------------------------------------- /praxis/layers/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/__init__.py -------------------------------------------------------------------------------- /praxis/layers/activations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/activations.py -------------------------------------------------------------------------------- /praxis/layers/activations_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/activations_test.py -------------------------------------------------------------------------------- /praxis/layers/adapters.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/adapters.py -------------------------------------------------------------------------------- /praxis/layers/adapters_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/adapters_test.py -------------------------------------------------------------------------------- /praxis/layers/attentions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/attentions.py -------------------------------------------------------------------------------- /praxis/layers/attentions_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/attentions_test.py -------------------------------------------------------------------------------- /praxis/layers/augmentations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/augmentations.py -------------------------------------------------------------------------------- /praxis/layers/augmentations_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/augmentations_test.py -------------------------------------------------------------------------------- /praxis/layers/base_ops.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/base_ops.py -------------------------------------------------------------------------------- /praxis/layers/bregman.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/bregman.py -------------------------------------------------------------------------------- /praxis/layers/bregman_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/bregman_test.py -------------------------------------------------------------------------------- /praxis/layers/chain/BUILD: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/chain/BUILD -------------------------------------------------------------------------------- /praxis/layers/chain/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/chain/__init__.py -------------------------------------------------------------------------------- /praxis/layers/chain/chain.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/chain/chain.py -------------------------------------------------------------------------------- /praxis/layers/chain/chain_extensions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/chain/chain_extensions.py -------------------------------------------------------------------------------- /praxis/layers/chain/chain_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/chain/chain_test.py -------------------------------------------------------------------------------- /praxis/layers/checkpoint_policy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/checkpoint_policy.py -------------------------------------------------------------------------------- /praxis/layers/chunk.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/chunk.py -------------------------------------------------------------------------------- /praxis/layers/chunk_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/chunk_test.py -------------------------------------------------------------------------------- /praxis/layers/conformers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/conformers.py -------------------------------------------------------------------------------- /praxis/layers/conformers_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/conformers_test.py -------------------------------------------------------------------------------- /praxis/layers/convolutions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/convolutions.py -------------------------------------------------------------------------------- /praxis/layers/convolutions_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/convolutions_test.py -------------------------------------------------------------------------------- /praxis/layers/ctc_objectives.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/ctc_objectives.py -------------------------------------------------------------------------------- /praxis/layers/ctc_objectives_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/ctc_objectives_test.py -------------------------------------------------------------------------------- /praxis/layers/einsum.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/einsum.py -------------------------------------------------------------------------------- /praxis/layers/einsum_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/einsum_test.py -------------------------------------------------------------------------------- /praxis/layers/embedding_softmax.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/embedding_softmax.py -------------------------------------------------------------------------------- /praxis/layers/embedding_softmax_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/embedding_softmax_test.py -------------------------------------------------------------------------------- /praxis/layers/flax_adapter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/flax_adapter.py -------------------------------------------------------------------------------- /praxis/layers/flax_adapter_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/flax_adapter_test.py -------------------------------------------------------------------------------- /praxis/layers/frnn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/frnn.py -------------------------------------------------------------------------------- /praxis/layers/frnn_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/frnn_test.py -------------------------------------------------------------------------------- /praxis/layers/glam.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/glam.py -------------------------------------------------------------------------------- /praxis/layers/gpu_fast_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/gpu_fast_attention.py -------------------------------------------------------------------------------- /praxis/layers/grok.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/grok.py -------------------------------------------------------------------------------- /praxis/layers/grouped_query_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/grouped_query_attention.py -------------------------------------------------------------------------------- /praxis/layers/grouped_query_attention_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/grouped_query_attention_test.py -------------------------------------------------------------------------------- /praxis/layers/injection/BUILD: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/injection/BUILD -------------------------------------------------------------------------------- /praxis/layers/injection/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/injection/__init__.py -------------------------------------------------------------------------------- /praxis/layers/injection/fp8_nvidia_gpu.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/injection/fp8_nvidia_gpu.py -------------------------------------------------------------------------------- /praxis/layers/injection/fp8_nvidia_gpu_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/injection/fp8_nvidia_gpu_test.py -------------------------------------------------------------------------------- /praxis/layers/linears.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/linears.py -------------------------------------------------------------------------------- /praxis/layers/linears_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/linears_test.py -------------------------------------------------------------------------------- /praxis/layers/losses.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/losses.py -------------------------------------------------------------------------------- /praxis/layers/mobilenet.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/mobilenet.py -------------------------------------------------------------------------------- /praxis/layers/models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/models.py -------------------------------------------------------------------------------- /praxis/layers/models_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/models_test.py -------------------------------------------------------------------------------- /praxis/layers/multi_query_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/multi_query_attention.py -------------------------------------------------------------------------------- /praxis/layers/multi_query_attention_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/multi_query_attention_test.py -------------------------------------------------------------------------------- /praxis/layers/ngrammer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/ngrammer.py -------------------------------------------------------------------------------- /praxis/layers/ngrammer_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/ngrammer_test.py -------------------------------------------------------------------------------- /praxis/layers/normalizations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/normalizations.py -------------------------------------------------------------------------------- /praxis/layers/normalizations_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/normalizations_test.py -------------------------------------------------------------------------------- /praxis/layers/pipeline.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/pipeline.py -------------------------------------------------------------------------------- /praxis/layers/pipeline_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/pipeline_test.py -------------------------------------------------------------------------------- /praxis/layers/poolings.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/poolings.py -------------------------------------------------------------------------------- /praxis/layers/poolings_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/poolings_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/BUILD: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/BUILD -------------------------------------------------------------------------------- /praxis/layers/quantization/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/__init__.py -------------------------------------------------------------------------------- /praxis/layers/quantization/attentions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/attentions.py -------------------------------------------------------------------------------- /praxis/layers/quantization/attentions_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/attentions_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/automl_select.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/automl_select.py -------------------------------------------------------------------------------- /praxis/layers/quantization/automl_select_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/automl_select_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/conformers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/conformers.py -------------------------------------------------------------------------------- /praxis/layers/quantization/conformers_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/conformers_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/convolutions.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/convolutions.py -------------------------------------------------------------------------------- /praxis/layers/quantization/convolutions_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/convolutions_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/einsum.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/einsum.py -------------------------------------------------------------------------------- /praxis/layers/quantization/einsum_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/einsum_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/embedding_softmax.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/embedding_softmax.py -------------------------------------------------------------------------------- /praxis/layers/quantization/embedding_softmax_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/embedding_softmax_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/linears.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/linears.py -------------------------------------------------------------------------------- /praxis/layers/quantization/linears_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/linears_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/multi_query_attention.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/multi_query_attention.py -------------------------------------------------------------------------------- /praxis/layers/quantization/multi_query_attention_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/multi_query_attention_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/ngrammer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/ngrammer.py -------------------------------------------------------------------------------- /praxis/layers/quantization/ngrammer_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/ngrammer_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/operations.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/operations.py -------------------------------------------------------------------------------- /praxis/layers/quantization/operations_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/operations_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/optimization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/optimization.py -------------------------------------------------------------------------------- /praxis/layers/quantization/optimization_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/optimization_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/overflow_check.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/overflow_check.py -------------------------------------------------------------------------------- /praxis/layers/quantization/overflow_check_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/overflow_check_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/quantization_hparams.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/quantization_hparams.py -------------------------------------------------------------------------------- /praxis/layers/quantization/quantization_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/quantization_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/quantize.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/quantize.py -------------------------------------------------------------------------------- /praxis/layers/quantization/quantize_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/quantize_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/quantizer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/quantizer.py -------------------------------------------------------------------------------- /praxis/layers/quantization/quantizer_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/quantizer_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/searchable.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/searchable.py -------------------------------------------------------------------------------- /praxis/layers/quantization/searchable_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/searchable_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/sparsity/BUILD: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/sparsity/BUILD -------------------------------------------------------------------------------- /praxis/layers/quantization/sparsity/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/sparsity/__init__.py -------------------------------------------------------------------------------- /praxis/layers/quantization/sparsity/attentions_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/sparsity/attentions_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/sparsity/linears_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/sparsity/linears_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/sparsity/sparsifier.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/sparsity/sparsifier.py -------------------------------------------------------------------------------- /praxis/layers/quantization/sparsity/sparsifier_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/sparsity/sparsifier_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/sparsity/sparsity.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/sparsity/sparsity.py -------------------------------------------------------------------------------- /praxis/layers/quantization/sparsity/sparsity_hparams.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/sparsity/sparsity_hparams.py -------------------------------------------------------------------------------- /praxis/layers/quantization/sparsity/sparsity_modes.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/sparsity/sparsity_modes.py -------------------------------------------------------------------------------- /praxis/layers/quantization/sparsity/sparsity_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/sparsity/sparsity_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/tests/attention_projection_aqt_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/tests/attention_projection_aqt_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/tests/attention_projection_fq_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/tests/attention_projection_fq_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/tests/attention_projection_ptq_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/tests/attention_projection_ptq_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/tests/combined_qkv_projection_aqt_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/tests/combined_qkv_projection_aqt_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/tests/combined_qkv_projection_fq_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/tests/combined_qkv_projection_fq_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/tests/combined_qkv_projection_ptq_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/tests/combined_qkv_projection_ptq_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/tests/dotproduct_attention_aqt_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/tests/dotproduct_attention_aqt_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/tests/linear_ptq_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/tests/linear_ptq_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/tests/one_headed_attention_projection_aqt_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/tests/one_headed_attention_projection_aqt_test.py -------------------------------------------------------------------------------- /praxis/layers/quantization/tests/test_util.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/tests/test_util.py -------------------------------------------------------------------------------- /praxis/layers/quantization/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/utils.py -------------------------------------------------------------------------------- /praxis/layers/quantization/utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantization/utils_test.py -------------------------------------------------------------------------------- /praxis/layers/quantizer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantizer.py -------------------------------------------------------------------------------- /praxis/layers/quantizer_objectives.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantizer_objectives.py -------------------------------------------------------------------------------- /praxis/layers/quantizer_objectives_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantizer_objectives_test.py -------------------------------------------------------------------------------- /praxis/layers/quantizer_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/quantizer_test.py -------------------------------------------------------------------------------- /praxis/layers/repeats.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/repeats.py -------------------------------------------------------------------------------- /praxis/layers/repeats_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/repeats_test.py -------------------------------------------------------------------------------- /praxis/layers/resnets.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/resnets.py -------------------------------------------------------------------------------- /praxis/layers/rnn_cell.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/rnn_cell.py -------------------------------------------------------------------------------- /praxis/layers/rnn_cell_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/rnn_cell_test.py -------------------------------------------------------------------------------- /praxis/layers/searchable.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/searchable.py -------------------------------------------------------------------------------- /praxis/layers/searchable_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/searchable_test.py -------------------------------------------------------------------------------- /praxis/layers/sedd/sedd.ipynb: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/sedd/sedd.ipynb -------------------------------------------------------------------------------- /praxis/layers/sequential.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/sequential.py -------------------------------------------------------------------------------- /praxis/layers/sequential_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/sequential_test.py -------------------------------------------------------------------------------- /praxis/layers/sharding.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/sharding.py -------------------------------------------------------------------------------- /praxis/layers/shared_layers_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/shared_layers_test.py -------------------------------------------------------------------------------- /praxis/layers/spectrum_augmenter.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/spectrum_augmenter.py -------------------------------------------------------------------------------- /praxis/layers/spectrum_augmenter_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/spectrum_augmenter_test.py -------------------------------------------------------------------------------- /praxis/layers/ssm.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/ssm.py -------------------------------------------------------------------------------- /praxis/layers/ssm_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/ssm_test.py -------------------------------------------------------------------------------- /praxis/layers/ssm_transformers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/ssm_transformers.py -------------------------------------------------------------------------------- /praxis/layers/ssm_transformers_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/ssm_transformers_test.py -------------------------------------------------------------------------------- /praxis/layers/stats.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/stats.py -------------------------------------------------------------------------------- /praxis/layers/stats_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/stats_test.py -------------------------------------------------------------------------------- /praxis/layers/stochastics.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/stochastics.py -------------------------------------------------------------------------------- /praxis/layers/stochastics_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/stochastics_test.py -------------------------------------------------------------------------------- /praxis/layers/test_layers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/test_layers.py -------------------------------------------------------------------------------- /praxis/layers/transformer_models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/transformer_models.py -------------------------------------------------------------------------------- /praxis/layers/transformer_models_encoder_decoder_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/transformer_models_encoder_decoder_test.py -------------------------------------------------------------------------------- /praxis/layers/transformer_models_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/transformer_models_test.py -------------------------------------------------------------------------------- /praxis/layers/transformers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/transformers.py -------------------------------------------------------------------------------- /praxis/layers/transformers_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/transformers_test.py -------------------------------------------------------------------------------- /praxis/layers/vanillanets.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/vanillanets.py -------------------------------------------------------------------------------- /praxis/layers/vanillanets_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/vanillanets_test.py -------------------------------------------------------------------------------- /praxis/layers/video/BUILD: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/video/BUILD -------------------------------------------------------------------------------- /praxis/layers/video/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/video/__init__.py -------------------------------------------------------------------------------- /praxis/layers/video/enc_dec_3dcnn.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/video/enc_dec_3dcnn.py -------------------------------------------------------------------------------- /praxis/layers/video/enc_dec_3dcnn_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/video/enc_dec_3dcnn_test.py -------------------------------------------------------------------------------- /praxis/layers/video/losses.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/video/losses.py -------------------------------------------------------------------------------- /praxis/layers/video/losses_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/video/losses_test.py -------------------------------------------------------------------------------- /praxis/layers/video/quantizer.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/video/quantizer.py -------------------------------------------------------------------------------- /praxis/layers/video/quantizer_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/video/quantizer_test.py -------------------------------------------------------------------------------- /praxis/layers/video/vqvae.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/video/vqvae.py -------------------------------------------------------------------------------- /praxis/layers/video/vqvae_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/video/vqvae_test.py -------------------------------------------------------------------------------- /praxis/layers/vits.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/vits.py -------------------------------------------------------------------------------- /praxis/layers/vits_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/layers/vits_test.py -------------------------------------------------------------------------------- /praxis/lazy_loader.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/lazy_loader.py -------------------------------------------------------------------------------- /praxis/lingvo_lib.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/lingvo_lib.py -------------------------------------------------------------------------------- /praxis/metric_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/metric_utils.py -------------------------------------------------------------------------------- /praxis/optimizer_prefix_vectorization.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/optimizer_prefix_vectorization.py -------------------------------------------------------------------------------- /praxis/optimizer_prefix_vectorization_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/optimizer_prefix_vectorization_test.py -------------------------------------------------------------------------------- /praxis/optimizers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/optimizers.py -------------------------------------------------------------------------------- /praxis/optimizers_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/optimizers_test.py -------------------------------------------------------------------------------- /praxis/pax_fiddle.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pax_fiddle.py -------------------------------------------------------------------------------- /praxis/pax_fiddle_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pax_fiddle_test.py -------------------------------------------------------------------------------- /praxis/pip_package/build_pip_pkg.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pip_package/build_pip_pkg.sh -------------------------------------------------------------------------------- /praxis/pip_package/cloudbuild-postsubmit.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pip_package/cloudbuild-postsubmit.yaml -------------------------------------------------------------------------------- /praxis/pip_package/cloudbuild-presubmit.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pip_package/cloudbuild-presubmit.yaml -------------------------------------------------------------------------------- /praxis/pip_package/cloudbuild-release.yaml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pip_package/cloudbuild-release.yaml -------------------------------------------------------------------------------- /praxis/pip_package/postsubmit.Dockerfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pip_package/postsubmit.Dockerfile -------------------------------------------------------------------------------- /praxis/pip_package/prepare_release.sh: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pip_package/prepare_release.sh -------------------------------------------------------------------------------- /praxis/pip_package/presubmit.Dockerfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pip_package/presubmit.Dockerfile -------------------------------------------------------------------------------- /praxis/pip_package/release.Dockerfile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pip_package/release.Dockerfile -------------------------------------------------------------------------------- /praxis/pip_package/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pip_package/requirements.txt -------------------------------------------------------------------------------- /praxis/praxis.bzl: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/praxis.bzl -------------------------------------------------------------------------------- /praxis/py_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/py_utils.py -------------------------------------------------------------------------------- /praxis/py_utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/py_utils_test.py -------------------------------------------------------------------------------- /praxis/pytypes.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pytypes.py -------------------------------------------------------------------------------- /praxis/pytypes_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/pytypes_test.py -------------------------------------------------------------------------------- /praxis/sample_decode.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/sample_decode.py -------------------------------------------------------------------------------- /praxis/sample_decode_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/sample_decode_test.py -------------------------------------------------------------------------------- /praxis/schedules.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/schedules.py -------------------------------------------------------------------------------- /praxis/schedules_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/schedules_test.py -------------------------------------------------------------------------------- /praxis/test_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/test_utils.py -------------------------------------------------------------------------------- /praxis/token_samplers.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/token_samplers.py -------------------------------------------------------------------------------- /praxis/token_samplers_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/token_samplers_test.py -------------------------------------------------------------------------------- /praxis/trees.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/trees.py -------------------------------------------------------------------------------- /praxis/trees_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/praxis/trees_test.py -------------------------------------------------------------------------------- /requirements.in: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/requirements.in -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/praxis/HEAD/setup.py --------------------------------------------------------------------------------