ML/AI Framework Ecosystem

From High-Level Code to Hardware Execution

Comprehensive interactive guide mapping the complete journey of machine learning frameworks, from familiar Python APIs to optimized hardware execution across CPUs, GPUs, TPUs, and specialized AI chips.

⚠️ Accuracy Notice

This diagram aims to provide a comprehensive overview of the ML/AI framework ecosystem. However, the landscape evolves rapidly and some connections or details might be outdated or incomplete. Please verify information independently and refer to official documentation for the most current details.

ML Framework Ecosystem Flowchart

flowchart TD %% Layer 1: Frameworks subgraph L1["🎯 Layer 1: High-Level Frameworks"] direction TB PT["🔥 PyTorch
nn.Module"] TF["🟠 TensorFlow
Graph/Eager"] JAX["🔵 JAX
NumPy-style"] end %% Layer 2: Bridges / Frontends subgraph L2["🔗 Layer 2: Bridges / Frontends"] direction TB subgraph L2A["PyTorch Bridges"] direction TB TC["📦 torch.compile"] PXLA["📦 torch_xla"] TMLIR["📦 torch-mlir"] PTONNX["📦 torch.onnx"] end subgraph L2B["TensorFlow Bridges"] direction TB TFXLA["📦 tf.function (+XLA)"] TFMLIR["📦 TF-MLIR"] TF2TFL["📦 TFLite Converter (TF)"] TFONNX["📦 tf2onnx"] TF2TRT["📦 TF-TensorRT"] end subgraph L2C["JAX Bridges"] direction TB JAXJIT["📦 jax.jit"] JAX2ONNX["📦 jax2onnx"] end end %% Layer 3: IRs & Exchange subgraph L3["📋 Layer 3: IRs & Exchange Formats"] direction TB FX["📋 FX Graph (PyTorch)"] STABLEHLO["📦 StableHLO (OpenXLA)"] MLIR["🏗️ MLIR (TOSA/Linalg/etc.)"] ONNX["📦 ONNX (Exchange)"] end %% Converters subgraph L3C["🔄 Converters"] direction TB ONNXTFL["ONNX → TFLite"] ONNX2TRT["ONNX → TensorRT"] ONNX2COREML["ONNX → Core ML"] TORCHTRT["Torch-TensorRT"] end %% Layer 4: Compilers & Runtimes subgraph L4["⚡ Layer 4: Compilers & Runtimes"] direction TB IND["🔧 TorchInductor"] XLA["⚡ OpenXLA (XLA)"] TVM["🎯 TVM"] IREE["📦 IREE"] TRITON["🔺 Triton Kernels"] ORT["▶️ ONNX Runtime"] TFLITE["▶️ TFLite Interpreter"] TRT["▶️ TensorRT Engine"] COREMLRT["▶️ Core ML Runtime"] NNAPIRT["▶️ NNAPI Runtime"] end %% Layer 5: Hardware subgraph L5["🖥️ Layer 5: Hardware Targets"] direction TB CPU["💻 CPU"] GPU["🎮 GPU"] TPU["🧠 TPU"] DSP["📶 DSP/Hexagon"] ANE["🍎 Apple Neural Engine"] NPU["📱 Mobile NPU"] end %% Ordering L1 --> L2 --> L3 --> L3C --> L4 --> L5 %% PyTorch PT --> TC --> FX --> IND PT --> PXLA --> STABLEHLO PT --> TMLIR --> MLIR PT --> PTONNX --> ONNX PT -.-> TORCHTRT %% TensorFlow TF --> TFXLA --> STABLEHLO TF --> TFMLIR --> MLIR TF --> TF2TFL --> TFLITE TF --> TFONNX --> ONNX TF --> TF2TRT --> TRT %% JAX JAX --> JAXJIT --> STABLEHLO JAX --> JAX2ONNX --> ONNX %% IRs STABLEHLO --> XLA MLIR --> IREE ONNX --> TVM ONNX --> ORT %% optional via onnx-mlir ONNX -.-> MLIR %% Converters ONNX --> ONNXTFL --> TFLITE ONNX --> ONNX2TRT --> TRT ONNX --> ONNX2COREML --> COREMLRT TORCHTRT --> TRT %% Compilers/Runtimes IND --> TRITON --> GPU IND --> CPU XLA --> TPU XLA --> GPU XLA --> CPU IREE --> CPU IREE --> GPU IREE --> NNAPIRT TVM --> CPU TVM --> GPU TVM -.-> NNAPIRT ORT --> CPU ORT --> GPU ORT --> NNAPIRT ORT --> COREMLRT TFLITE --> CPU TFLITE --> GPU TFLITE --> NNAPIRT TFLITE --> DSP TFLITE --> COREMLRT TRT --> GPU COREMLRT --> ANE COREMLRT --> GPU COREMLRT --> CPU NNAPIRT --> NPU NNAPIRT --> CPU %% Styling classDef layer fill:#f8f9fa,stroke:#343a40,stroke-width:3px classDef framework fill:#e1f5fe,stroke:#0277bd,color:#000 classDef ir fill:#fff8e1,stroke:#ff8f00,color:#000 classDef converter fill:#fff3e0,stroke:#f57c00,color:#000 classDef compiler fill:#e8eaf6,stroke:#3f51b5,color:#000 classDef runtime fill:#ede7f6,stroke:#5e35b1,color:#000 classDef hardware fill:#e8f5e8,stroke:#2e7d32,color:#000 class L1,L2,L3,L3C,L4,L5 layer class PT,TF,JAX framework class FX,ONNX,STABLEHLO,MLIR ir class ONNXTFL,ONNX2TRT,ONNX2COREML,TORCHTRT converter class IND,XLA,TVM,IREE,TRITON compiler class ORT,TFLITE,TRT,COREMLRT,NNAPIRT runtime class CPU,GPU,TPU,DSP,ANE,NPU hardware

5-Layer Architecture Overview

1

High-Level Frameworks

Developer-friendly interfaces where machine learning models are defined using familiar Python APIs and high-level abstractions.

Examples:

PyTorch (nn.Module), TensorFlow (Graph/Eager), JAX (NumPy-style)

2

Bridges & Frontends

Translation layers that convert framework-specific code into intermediate representations for cross-platform optimization.

Examples:

torch.compile, tf.function, jax.jit, torch_xla, tf2onnx

3

Intermediate Representations

Standardized formats that enable model portability and serve as common interfaces between different tools in the ecosystem.

Examples:

ONNX, StableHLO, MLIR, FX Graph

4

Compilers & Runtimes

Optimization engines that transform intermediate representations into efficient, hardware-specific executable code.

Examples:

TorchInductor, XLA, TVM, ONNX Runtime, TensorRT

5

Hardware Targets

Physical computing platforms where optimized model code executes, ranging from general-purpose processors to specialized AI accelerators.

Examples:

CPUs, GPUs, TPUs, Mobile NPUs, Apple Neural Engine

Key Insights

Multiple Pathways

Each framework offers multiple compilation paths to the same hardware, providing flexibility in optimization strategies.

ONNX as Bridge

ONNX serves as the universal interchange format, enabling cross-framework model deployment.

Hardware Specialization

Different deployment targets have optimized runtimes tailored to their specific constraints and capabilities.

Resources

PyTorch Documentation

TensorFlow Guide

ONNX Documentation

JAX Documentation

OpenXLA Project