Noureddine RAMDI / Keras 3: Multi-backend deep learning framework simplifying model development across JAX, TensorFlow, and PyTorch

Created Sun, 26 Apr 2026 09:31:26 +0000 Modified Sat, 23 May 2026 20:41:27 +0000
955 Words
github-stars machine learning deep learning python keras jax tensorflow pytorch

keras-team/keras

Keras 3 takes a different approach to deep learning frameworks by supporting multiple backends—JAX, TensorFlow, PyTorch, and OpenVINO—under a unified API. This design lets developers write models once and run them on the backend best suited for their hardware or task, sidestepping the usual framework lock-in that often complicates ML projects. The repo claims speedups ranging from 20% to 350% compared to other frameworks, making it worth a close look for practitioners concerned with performance and flexibility.

What Keras 3 offers as a multi-backend deep learning framework

At its core, Keras 3 is a deep learning framework implemented in Python that emphasizes accelerated model development, scalability, and broad hardware support. Unlike traditional single-backend frameworks, Keras 3 abstracts the backend layer to support multiple execution engines. Currently, it supports three main training and inference backends: JAX, TensorFlow, and PyTorch. Additionally, it supports OpenVINO for inference-only workloads.

This multi-backend architecture is designed to be largely backward compatible with tf.keras, allowing users to migrate existing codebases with minimal changes. The framework supports typical deep learning tasks: building layers, models, training loops, and custom components. The key is that these components can run transparently on any supported backend.

Under the hood, Keras 3 achieves this by leveraging backend-specific implementations that conform to a common API. This architectural choice enables the same model code to run on CPUs, GPUs, TPUs, or specialized accelerators without modification. The framework also supports scaling from laptops to large distributed clusters.

The repo’s design encourages writing future-proof machine learning code that avoids framework lock-in. This is important because the deep learning landscape evolves rapidly, and being tied to a single backend can limit flexibility and performance optimizations.

Why Keras 3’s multi-backend architecture stands out

The standout feature of Keras 3 is its ability to unify three major ML frameworks (JAX, TensorFlow, PyTorch) under a consistent, user-friendly API. This is no small feat given the differences in execution models, APIs, and ecosystem tooling among these frameworks.

One of the tradeoffs is the complexity of maintaining compatibility and feature parity across backends. Some backend-specific features or optimizations may not be available universally, which can limit certain advanced use cases. For example, OpenVINO is supported only for inference, not training.

Despite this complexity, the codebase is surprisingly clean and modular. The project uses layered abstractions to separate backend-specific logic from the core API. This separation helps maintain code quality and makes it easier to add new backends or extend functionality.

Performance-wise, the repo claims speedups between 20% and 350% compared to other frameworks, depending on the backend and workload. This suggests that users can optimize execution by choosing the most appropriate backend for their task or hardware. For instance, JAX may offer advantages in research settings with its functional programming style and just-in-time compilation, while TensorFlow or PyTorch may be preferred for production deployment.

The backward compatibility with tf.keras is a practical benefit for teams migrating large codebases. It minimizes the friction of adopting Keras 3 while enabling incremental migration and experimentation with different backends.

Quick start with Keras 3

The installation process is straightforward and documented clearly in the repo. You install Keras 3 itself via pip, then install one or more backend packages as needed.

# Install Keras 3 package
pip install keras --upgrade

# Install backend package(s) of your choice
# Options include tensorflow, jax, or torch
# OpenVINO backend is available for inference only

For those interested in local development or contributing, the repo supports Linux and macOS natively, with Windows users advised to use WSL2.

# Install dependencies for local development
pip install -r requirements.txt

# Install Keras locally from the root directory
python pip_build.py --install

# When updating public APIs, regenerate API files
./shell/api_gen.sh

This setup balances ease of use for general users with flexibility for developers who want to dive into the source.

Verdict: who should consider Keras 3?

Keras 3 is relevant for deep learning practitioners who want to reduce framework lock-in and experiment across multiple backends without rewriting their models. Its multi-backend architecture makes it a solid choice for research teams that value flexibility and want to optimize performance by selecting the best backend per workload.

The framework also suits production environments that may need to scale models across diverse hardware or transition between frameworks over time. The backward compatibility with tf.keras lowers the barrier for existing TensorFlow users.

That said, the multi-backend approach introduces some tradeoffs. Not all backend features may be available uniformly, and the added abstraction can increase complexity. OpenVINO’s inference-only support is a limitation for those seeking unified training and inference workflows on that backend.

Overall, Keras 3’s approach to framework agnosticism is worth understanding even if you don’t adopt it immediately. It represents a practical step toward more portable and performant deep learning code in a fragmented ecosystem.

→ GitHub Repo: keras-team/keras ⭐ 64,041 · Python