├── .github └── workflows │ ├── publish.yml │ └── unittests.yml ├── CITATION.cff ├── CONTRIBUTING.md ├── LICENSE ├── MANIFEST.in ├── README.md ├── docs ├── .gitignore ├── Makefile ├── api.rst ├── conf.py ├── examples │ ├── approx_metrics.rst │ ├── flax_integration.rst │ ├── segmentation.rst │ └── t5x.rst ├── index.rst └── references.bib ├── examples ├── __init__.py ├── approx_metrics │ ├── README.md │ ├── __init__.py │ ├── main.py │ └── main_test.py ├── flax_integration │ ├── README.md │ ├── __init__.py │ ├── main.py │ └── main_test.py ├── segmentation │ ├── README.md │ ├── __init__.py │ ├── main.py │ └── main_test.py └── t5x │ ├── README.md │ ├── models.py │ ├── models_test.py │ ├── msmarco.gin │ ├── requirements.txt │ ├── tasks.py │ └── tasks_test.py ├── rax ├── __init__.py ├── _src │ ├── __init__.py │ ├── lambdaweights.py │ ├── lambdaweights_test.py │ ├── losses.py │ ├── losses_test.py │ ├── metrics.py │ ├── metrics_test.py │ ├── segment_utils.py │ ├── segment_utils_test.py │ ├── t12n.py │ ├── t12n_test.py │ ├── types.py │ ├── utils.py │ └── utils_test.py ├── types.py └── utils.py ├── readthedocs.yml ├── requirements ├── requirements-docs.txt ├── requirements-examples.txt ├── requirements-test.txt └── requirements.txt └── setup.py /.github/workflows/publish.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/.github/workflows/publish.yml -------------------------------------------------------------------------------- /.github/workflows/unittests.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/.github/workflows/unittests.yml -------------------------------------------------------------------------------- /CITATION.cff: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/CITATION.cff -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/CONTRIBUTING.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/LICENSE -------------------------------------------------------------------------------- /MANIFEST.in: -------------------------------------------------------------------------------- 1 | include requirements/* 2 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/README.md -------------------------------------------------------------------------------- /docs/.gitignore: -------------------------------------------------------------------------------- 1 | _build 2 | -------------------------------------------------------------------------------- /docs/Makefile: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/docs/Makefile -------------------------------------------------------------------------------- /docs/api.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/docs/api.rst -------------------------------------------------------------------------------- /docs/conf.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/docs/conf.py -------------------------------------------------------------------------------- /docs/examples/approx_metrics.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/docs/examples/approx_metrics.rst -------------------------------------------------------------------------------- /docs/examples/flax_integration.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/docs/examples/flax_integration.rst -------------------------------------------------------------------------------- /docs/examples/segmentation.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/docs/examples/segmentation.rst -------------------------------------------------------------------------------- /docs/examples/t5x.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/docs/examples/t5x.rst -------------------------------------------------------------------------------- /docs/index.rst: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/docs/index.rst -------------------------------------------------------------------------------- /docs/references.bib: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/docs/references.bib -------------------------------------------------------------------------------- /examples/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/__init__.py -------------------------------------------------------------------------------- /examples/approx_metrics/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/approx_metrics/README.md -------------------------------------------------------------------------------- /examples/approx_metrics/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/approx_metrics/__init__.py -------------------------------------------------------------------------------- /examples/approx_metrics/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/approx_metrics/main.py -------------------------------------------------------------------------------- /examples/approx_metrics/main_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/approx_metrics/main_test.py -------------------------------------------------------------------------------- /examples/flax_integration/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/flax_integration/README.md -------------------------------------------------------------------------------- /examples/flax_integration/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/flax_integration/__init__.py -------------------------------------------------------------------------------- /examples/flax_integration/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/flax_integration/main.py -------------------------------------------------------------------------------- /examples/flax_integration/main_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/flax_integration/main_test.py -------------------------------------------------------------------------------- /examples/segmentation/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/segmentation/README.md -------------------------------------------------------------------------------- /examples/segmentation/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/segmentation/__init__.py -------------------------------------------------------------------------------- /examples/segmentation/main.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/segmentation/main.py -------------------------------------------------------------------------------- /examples/segmentation/main_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/segmentation/main_test.py -------------------------------------------------------------------------------- /examples/t5x/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/t5x/README.md -------------------------------------------------------------------------------- /examples/t5x/models.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/t5x/models.py -------------------------------------------------------------------------------- /examples/t5x/models_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/t5x/models_test.py -------------------------------------------------------------------------------- /examples/t5x/msmarco.gin: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/t5x/msmarco.gin -------------------------------------------------------------------------------- /examples/t5x/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/t5x/requirements.txt -------------------------------------------------------------------------------- /examples/t5x/tasks.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/t5x/tasks.py -------------------------------------------------------------------------------- /examples/t5x/tasks_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/examples/t5x/tasks_test.py -------------------------------------------------------------------------------- /rax/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/__init__.py -------------------------------------------------------------------------------- /rax/_src/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/__init__.py -------------------------------------------------------------------------------- /rax/_src/lambdaweights.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/lambdaweights.py -------------------------------------------------------------------------------- /rax/_src/lambdaweights_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/lambdaweights_test.py -------------------------------------------------------------------------------- /rax/_src/losses.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/losses.py -------------------------------------------------------------------------------- /rax/_src/losses_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/losses_test.py -------------------------------------------------------------------------------- /rax/_src/metrics.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/metrics.py -------------------------------------------------------------------------------- /rax/_src/metrics_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/metrics_test.py -------------------------------------------------------------------------------- /rax/_src/segment_utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/segment_utils.py -------------------------------------------------------------------------------- /rax/_src/segment_utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/segment_utils_test.py -------------------------------------------------------------------------------- /rax/_src/t12n.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/t12n.py -------------------------------------------------------------------------------- /rax/_src/t12n_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/t12n_test.py -------------------------------------------------------------------------------- /rax/_src/types.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/types.py -------------------------------------------------------------------------------- /rax/_src/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/utils.py -------------------------------------------------------------------------------- /rax/_src/utils_test.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/_src/utils_test.py -------------------------------------------------------------------------------- /rax/types.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/types.py -------------------------------------------------------------------------------- /rax/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/rax/utils.py -------------------------------------------------------------------------------- /readthedocs.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/readthedocs.yml -------------------------------------------------------------------------------- /requirements/requirements-docs.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/requirements/requirements-docs.txt -------------------------------------------------------------------------------- /requirements/requirements-examples.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/requirements/requirements-examples.txt -------------------------------------------------------------------------------- /requirements/requirements-test.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/requirements/requirements-test.txt -------------------------------------------------------------------------------- /requirements/requirements.txt: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/requirements/requirements.txt -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google/rax/HEAD/setup.py --------------------------------------------------------------------------------