├── .github └── workflows │ └── ci.yml ├── CHANGELOG.md ├── CONTRIBUTING.md ├── LICENSE ├── README.md ├── benchmarks ├── README.md └── benchmark.py ├── docs └── coverage.md ├── examples ├── jax_time_batch.py └── tensorflow_time_batch.py ├── jax-stubs ├── jax-stubs │ ├── __init__.py │ ├── __init__.pyi │ ├── nn │ │ └── __init__.pyi │ └── numpy │ │ └── __init__.pyi └── setup.py ├── setup.py ├── tensor_annotations ├── __init__.py ├── axes.py ├── jax.py ├── jax.pyi ├── library_stubs │ └── third_party │ │ └── py │ │ └── numpy │ │ └── __init__.pyi ├── numpy.py ├── numpy.pyi ├── py.typed ├── templates │ ├── jax.pyi │ ├── jax_nn.pyi │ ├── jax_numpy.pyi │ ├── jax_tensors.pyi │ ├── numpy.pyi │ ├── numpy_tensors.pyi │ ├── tensorflow.pyi │ └── tensorflow_tensors.pyi ├── tensorflow.py ├── tensorflow.pyi ├── tests │ ├── jax.py │ ├── numpy.py │ ├── pytype.py │ ├── templates.py │ ├── tensorflow.py │ └── utils.py └── tools │ ├── render_jax_library_template.py │ ├── render_numpy_library_template.py │ ├── render_tensor_template.py │ ├── render_tensorflow_library_template.py │ └── templates.py └── tensorflow-stubs ├── setup.py └── tensorflow-stubs ├── __init__.py ├── __init__.pyi ├── core └── util │ └── __init__.pyi ├── math └── __init__.pyi ├── python ├── __init__.pyi ├── ops │ └── __init__.pyi └── util │ └── __init__.pyi └── train └── __init__.pyi /.github/workflows/ci.yml: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/.github/workflows/ci.yml -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/CHANGELOG.md -------------------------------------------------------------------------------- /CONTRIBUTING.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/CONTRIBUTING.md -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/LICENSE -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/README.md -------------------------------------------------------------------------------- /benchmarks/README.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/benchmarks/README.md -------------------------------------------------------------------------------- /benchmarks/benchmark.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/benchmarks/benchmark.py -------------------------------------------------------------------------------- /docs/coverage.md: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/docs/coverage.md -------------------------------------------------------------------------------- /examples/jax_time_batch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/examples/jax_time_batch.py -------------------------------------------------------------------------------- /examples/tensorflow_time_batch.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/examples/tensorflow_time_batch.py -------------------------------------------------------------------------------- /jax-stubs/jax-stubs/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /jax-stubs/jax-stubs/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/jax-stubs/jax-stubs/__init__.pyi -------------------------------------------------------------------------------- /jax-stubs/jax-stubs/nn/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/jax-stubs/jax-stubs/nn/__init__.pyi -------------------------------------------------------------------------------- /jax-stubs/jax-stubs/numpy/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/jax-stubs/jax-stubs/numpy/__init__.pyi -------------------------------------------------------------------------------- /jax-stubs/setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/jax-stubs/setup.py -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/setup.py -------------------------------------------------------------------------------- /tensor_annotations/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/__init__.py -------------------------------------------------------------------------------- /tensor_annotations/axes.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/axes.py -------------------------------------------------------------------------------- /tensor_annotations/jax.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/jax.py -------------------------------------------------------------------------------- /tensor_annotations/jax.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/jax.pyi -------------------------------------------------------------------------------- /tensor_annotations/library_stubs/third_party/py/numpy/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/library_stubs/third_party/py/numpy/__init__.pyi -------------------------------------------------------------------------------- /tensor_annotations/numpy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/numpy.py -------------------------------------------------------------------------------- /tensor_annotations/numpy.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/numpy.pyi -------------------------------------------------------------------------------- /tensor_annotations/py.typed: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/py.typed -------------------------------------------------------------------------------- /tensor_annotations/templates/jax.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/templates/jax.pyi -------------------------------------------------------------------------------- /tensor_annotations/templates/jax_nn.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/templates/jax_nn.pyi -------------------------------------------------------------------------------- /tensor_annotations/templates/jax_numpy.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/templates/jax_numpy.pyi -------------------------------------------------------------------------------- /tensor_annotations/templates/jax_tensors.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/templates/jax_tensors.pyi -------------------------------------------------------------------------------- /tensor_annotations/templates/numpy.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/templates/numpy.pyi -------------------------------------------------------------------------------- /tensor_annotations/templates/numpy_tensors.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/templates/numpy_tensors.pyi -------------------------------------------------------------------------------- /tensor_annotations/templates/tensorflow.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/templates/tensorflow.pyi -------------------------------------------------------------------------------- /tensor_annotations/templates/tensorflow_tensors.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/templates/tensorflow_tensors.pyi -------------------------------------------------------------------------------- /tensor_annotations/tensorflow.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tensorflow.py -------------------------------------------------------------------------------- /tensor_annotations/tensorflow.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tensorflow.pyi -------------------------------------------------------------------------------- /tensor_annotations/tests/jax.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tests/jax.py -------------------------------------------------------------------------------- /tensor_annotations/tests/numpy.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tests/numpy.py -------------------------------------------------------------------------------- /tensor_annotations/tests/pytype.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tests/pytype.py -------------------------------------------------------------------------------- /tensor_annotations/tests/templates.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tests/templates.py -------------------------------------------------------------------------------- /tensor_annotations/tests/tensorflow.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tests/tensorflow.py -------------------------------------------------------------------------------- /tensor_annotations/tests/utils.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tests/utils.py -------------------------------------------------------------------------------- /tensor_annotations/tools/render_jax_library_template.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tools/render_jax_library_template.py -------------------------------------------------------------------------------- /tensor_annotations/tools/render_numpy_library_template.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tools/render_numpy_library_template.py -------------------------------------------------------------------------------- /tensor_annotations/tools/render_tensor_template.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tools/render_tensor_template.py -------------------------------------------------------------------------------- /tensor_annotations/tools/render_tensorflow_library_template.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tools/render_tensorflow_library_template.py -------------------------------------------------------------------------------- /tensor_annotations/tools/templates.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensor_annotations/tools/templates.py -------------------------------------------------------------------------------- /tensorflow-stubs/setup.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensorflow-stubs/setup.py -------------------------------------------------------------------------------- /tensorflow-stubs/tensorflow-stubs/__init__.py: -------------------------------------------------------------------------------- 1 | -------------------------------------------------------------------------------- /tensorflow-stubs/tensorflow-stubs/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensorflow-stubs/tensorflow-stubs/__init__.pyi -------------------------------------------------------------------------------- /tensorflow-stubs/tensorflow-stubs/core/util/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensorflow-stubs/tensorflow-stubs/core/util/__init__.pyi -------------------------------------------------------------------------------- /tensorflow-stubs/tensorflow-stubs/math/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensorflow-stubs/tensorflow-stubs/math/__init__.pyi -------------------------------------------------------------------------------- /tensorflow-stubs/tensorflow-stubs/python/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensorflow-stubs/tensorflow-stubs/python/__init__.pyi -------------------------------------------------------------------------------- /tensorflow-stubs/tensorflow-stubs/python/ops/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensorflow-stubs/tensorflow-stubs/python/ops/__init__.pyi -------------------------------------------------------------------------------- /tensorflow-stubs/tensorflow-stubs/python/util/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensorflow-stubs/tensorflow-stubs/python/util/__init__.pyi -------------------------------------------------------------------------------- /tensorflow-stubs/tensorflow-stubs/train/__init__.pyi: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/google-deepmind/tensor_annotations/HEAD/tensorflow-stubs/tensorflow-stubs/train/__init__.pyi --------------------------------------------------------------------------------