| commit | c397258f87007840b8b66e66789400a9daef1c68 | [log] [tgz] |
|---|---|---|
| author | Scott Todd <scotttodd@google.com> | Fri Oct 13 15:18:33 2023 -0700 |
| committer | GitHub <noreply@github.com> | Fri Oct 13 15:18:33 2023 -0700 |
| tree | fd19aa9fab6ccd27a9c22e3af102d9421ba6f719 | |
| parent | 3f001b695cc6d271da02d97aa2b2641bc7baf3b4 [diff] [blame] |
Add pytorch_aot_simple sample Colab notebook using SHARK-Turbine. (#15166) Progress on https://github.com/openxla/iree/issues/15117 From [SHARK-Turbine](https://github.com/nod-ai/SHARK-Turbine)'s README: > *AOT Export*: For compiling one or more `nn.Module`s to compiled, deployment ready artifacts. This operates via both a [simple one-shot export API](https://github.com/nod-ai/SHARK-Turbine/blob/main/python/shark_turbine/aot/exporter.py) for simple models and an underlying [advanced API](https://github.com/nod-ai/SHARK-Turbine/blob/main/python/shark_turbine/aot/compiled_module.py) for complicated models and accessing the full features of the runtime. This notebook shows how to use the "simple one-shot export API", with a branch between in-session compilation using `export_output.compile()` and export to native tools using `export_output.save_mlir()`. Some portion of this notebook may be included as a code sample for https://github.com/openxla/iree/issues/15114. Preview URL for review: https://colab.research.google.com/github/scotttodd/iree/blob/pytorch-samples-2/samples/colab/pytorch_aot_simple.ipynb skip-ci: no-op
diff --git a/samples/colab/test_notebooks.py b/samples/colab/test_notebooks.py index 17ca31e..d441eec 100755 --- a/samples/colab/test_notebooks.py +++ b/samples/colab/test_notebooks.py
@@ -15,7 +15,8 @@ # tflite_runtime requires some deps ("version `GLIBC_2.29' not found") that # samples.Dockerfile does not currently include. "tflite_text_classification.ipynb", - # Requires Python 3.10+ in our Docker image. + # PyTorch notebooks using SHARK-Turbine require Python 3.10+ in Docker. + "pytorch_aot_simple.ipynb", "pytorch_jit.ipynb", ]