Keras

Keras 3 is a game changer: it is now multi-backend. You can write Keras code and run it on top of JAX, PyTorch, or TensorFlow.

When to Use

• Multi-Backend Deep Learning (Keras 3): Writing models that run interchangeably on JAX, PyTorch, or TensorFlow.
• Rapid Neural Network Prototyping: Developing vision, NLP, and tabular architectures with clean, intuitive APIs.
• Production Model Deployment: Exporting models to ONNX, TensorRT, or TensorFlow Lite.
• Custom Layers & Loss Functions: Subclassing keras.layers.Layer with backend-agnostic keras.ops.

Quick Start

import os
os.environ["KERAS_BACKEND"] = "torch" # Choose: jax, torch, or tensorflow
import keras
from keras import layers