TT-XLA leverages a PJRT interface to integrate JAX (and in the future other frameworks), TT-MLIR and Tenstorrent hardware. It supports ingestion of PyTorch models through PyTorch/XLA and JAX models via jit compile, providing a StableHLO (SHLO) graph to TT-MLIR compiler.

---

Quick Links

• Getting Started / How to Run a Model

---

What is This Repo?

The TT-XLA repository is primarily used to enable running PyTorch and JAX models on Tenstorrent's AI hardware. It's a backend integration between the JAX ecosystem and Tenstorrent's ML accelerators using the PJRT (Portable JAX Runtime) interface.

---

Current AI Framework Front End Projects

• TT-XLA

  • TT-XLA is the primary frontend for running PyTorch and JAX models. It leverages a PJRT interface to integrate JAX (and in the future other frameworks), TT-MLIR, and Tenstorrent hardware. It supports ingestion of JAX models via jit compile, providing StableHLO (SHLO) graph to TT-MLIR compiler. TT-XLA can be used for single and multi-chip projects.
  • See the TT-XLA docs pages for an overview and getting started guide.

• TT-Forge-ONNX

  • A TVM based graph compiler designed to optimize and transform computational graphs for deep learning models. Supports ingestion of ONNX, TensorFlow, PaddlePaddle and similar ML frameworks via TVM (TT-TVM). It also supports ingestion of PyTorch, however it is recommended that you use TT-XLA. TT-Forge-ONNX does not support multi-chip configurations; it is for single-chip projects only.
  • See the TT-Forge-ONNX docs pages for an overview and getting started guide.

• TT-Torch - (deprecated)

  • A MLIR-native, open-source, PyTorch 2.X and torch-mlir based front-end. It provides stableHLO (SHLO) graphs to TT-MLIR. Supports ingestion of PyTorch models via PT2.X compile and ONNX models via torch-mlir (ONNX->SHLO)
  • See the TT-Torch docs pages for an overview and getting started guide.

---

Related Tenstorrent Projects

• TT-XLA - (single and multi-chip) For use with PyTorch, JAX, and TensorFlow
• TT-Forge-ONNX - (single chip only) For use with ONNX and PaddlePaddle, it also runs PyTorch, however it is recommended to use TT-XLA for PyTorch
• TT-MLIR - Open source compiler framework for compiling and optimizing machine learning models for Tenstorrent hardware
• TT-Metal - Low-level programming model, enabling kernel development for Tenstorrent hardware
• TT-TVM - A compiler stack for deep learning systems designed to close the gap between the productivity-focused deep learning frameworks, and the performance and efficiency-focused hardware backends
• TT-Torch - (Deprecated) Previously for use with PyTorch. It is recommended that you use TT-XLA for PyTorch.

---

Tenstorrent Bounty Program Terms and Conditions

This repo is a part of Tenstorrent’s bounty program. If you are interested in helping to improve tt-forge, please make sure to read the Tenstorrent Bounty Program Terms and Conditions before heading to the issues tab. Look for the issues that are tagged with both “bounty” and difficulty level!