Skip to main content

Multi-backend Keras.

Project description

Keras Core: A new multi-backend Keras

Keras Core is a new multi-backend implementation of the Keras API, with support for TensorFlow, JAX, and PyTorch.

WARNING: At this time, this package is experimental. It has rough edges and not everything might work as expected. We are currently hard at work improving it.

Once ready, this package will become Keras 3.0 and subsume tf.keras.

Local installation

Keras Core is compatible with Linux and MacOS systems. To install a local development version:

  1. Install dependencies:
pip install -r requirements.txt
  1. Run installation command from the root directory.
python pip_build.py --install

You should also install your backend of choice: tensorflow, jax, or torch. Note that tensorflow is required for using certain Keras Core features: certain preprocessing layers as well as tf.data pipelines.

Configuring your backend

You can export the environment variable KERAS_BACKEND or you can edit your local config file at ~/.keras/keras.json to configure your backend. Available backend options are: "tensorflow", "jax", "torch". Example:

export KERAS_BACKEND="jax"

In Colab, you can do:

import os
os.environ["KERAS_BACKEND"] = "jax"

import keras_core as keras

Note: The backend must be configured before importing keras_core, and the backend cannot be changed after the package has been imported.

Backwards compatibility

Keras Core is intended to work as a drop-in replacement for tf.keras (when using the TensorFlow backend). Just take your existing tf.keras code, change the keras imports to keras_core, make sure that your calls to model.save() are using the up-to-date .keras format, and you're done.

If your tf.keras model does not include custom components, you can start running it on top of JAX or PyTorch immediately.

If it does include custom components (e.g. custom layers or a custom train_step()), it is usually possible to convert it to a backend-agnostic implementation in just a few minutes.

In addition, Keras models can consume datasets in any format, regardless of the backend you're using: you can train your models with your existing tf.data.Dataset pipelines or PyTorch DataLoaders.

Why use Keras Core?

  • Run your high-level Keras workflows on top of any framework -- benefiting at will from the advantages of each framework, e.g. the scalability and performance of JAX or the production ecosystem options of TensorFlow.
  • Write custom components (e.g. layers, models, metrics) that you can use in low-level workflows in any framework.
    • You can take a Keras model and train it in a training loop written from scratch in native TF, JAX, or PyTorch.
    • You can take a Keras model and use it as part of a PyTorch-native Module or as part of a JAX-native model function.
  • Make your ML code future-proof by avoiding framework lock-in.
  • As a PyTorch user: get access to power and usability of Keras, at last!
  • As a JAX user: get access to a fully-featured, battle-tested, well-documented modeling and training library.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

keras-core-0.1.7.tar.gz (681.4 kB view details)

Uploaded Source

Built Distribution

keras_core-0.1.7-py3-none-any.whl (950.8 kB view details)

Uploaded Python 3

File details

Details for the file keras-core-0.1.7.tar.gz.

File metadata

  • Download URL: keras-core-0.1.7.tar.gz
  • Upload date:
  • Size: 681.4 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.2 CPython/3.10.7

File hashes

Hashes for keras-core-0.1.7.tar.gz
Algorithm Hash digest
SHA256 a68580dc6910662026afdb5cda21e157f5163c6a77e86beae0e8d1de40234c07
MD5 f45ce3d520d357d73254b8aa402364da
BLAKE2b-256 99c58ad2a43f3ba78a1305bbf84a38fd0092207fd4abb69c37b7c88898cdb991

See more details on using hashes here.

File details

Details for the file keras_core-0.1.7-py3-none-any.whl.

File metadata

  • Download URL: keras_core-0.1.7-py3-none-any.whl
  • Upload date:
  • Size: 950.8 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? No
  • Uploaded via: twine/4.0.2 CPython/3.10.7

File hashes

Hashes for keras_core-0.1.7-py3-none-any.whl
Algorithm Hash digest
SHA256 01f9f322d4b7e6e0a64d9cc3de022144ddf82ff88a0ac2b68e49f1a67d4a7ff0
MD5 937f874183fe252bbc19f4145dff2963
BLAKE2b-256 95f7b8dcff937ea64f822f0d3fe8c6010793406b82d14467cd0e9eecea458a40

See more details on using hashes here.

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page