Top AI Repos — open-source AI, indexed and scored
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
Top AI Repos tracks AI repositories on GitHub and answers two different questions about each one: is it moving right now, and would you bet a product on it.
A High Level API for Deep Learning in JAX
| Date | Stars |
|---|---|
| 2026-07-31 | 473 |
| 2026-08-04 | 473 |
| 2026-08-06 | 473 |
Today
— stars today
This week
— stars this week
This month
— stars this month
Momentum
0.0
growth rate 0.00%/day
# Elegy
<!-- [](https://pypi.org/project/elegy/) -->
<!-- [](https://pypi.org/project/elegy/) -->
<!-- [](https://poets-ai.github.io/elegy/) -->
<!-- [](https://github.com/psf/black) -->
[](https://codecov.io/gh/poets-ai/elegy)
[](https://github.com/poets-ai/elegy/actions?query=workflow%3A%22GitHub+CI%22)
[](https://github.com/poets-ai/elegy/issues)
______________________________________________________________________
_A High Level API for Deep Learning in JAX_
#### Main Features
- 😀 **Easy-to-use**: Elegy provides a Keras-like high-level API that makes it very easy to use for most common tasks.
- 💪 **Flexible**: Elegy provides a Pytorch Lightning-like low-level API that offers maximum flexibility when needed.
- 🔌 **Compatible**: Elegy supports various frameworks and data sources including Flax & Haiku Modules, Optax Optimizers, TensorFlow Datasets, Pytorch DataLoaders, and more.
<!-- - 🤷 **Agnostic**: Elegy supports various frameworks, including Flax, Haiku, and Optax on the high-level API, and it is 100% framework-agnostic on the low-level API. -->
Elegy is built on top of [Treex](https://github.com/cgarciae/treex) and [Treeo](https://github.com/cgarciae/treeo) and reexports their APIs for convenience.
[Getting Started](https://poets-ai.github.io/elegy/getting-started/high-level-api) | [Examples](/examples) | [Documentation](https://poets-ai.github.io/elegy)
## What is included?
* A `Model` class with an Estimator-like API.
* A `callbacks` module with common Keras callbacks.
**From Treex**
* A `Module` class.
* A `nn` module for with common layers.
* A `losses` module with common loss functions.
* A `metrics` module with common metrics.
## Installation
Install using pip:
```bash
pip install elegy
```
For Windows users, we recommend the Windows subsystem for Linux 2 [WSL2](https://docs.microsoft.com/es-es/windows/wsl/install-win10?redirectedfrom=MSDN) since [jax](https://github.com/google/jax/issues/438) does not support it yet.
## Quick Start: High-level API
Elegy's high-level API provides a straightforward interface you can use by implementing the following steps:
**1.** Define the architecture inside a `Module`:
```python
import jax
import elegy as eg
class MLP(eg.Module):
@eg.compact
def __call__(self, x):
x = eg.Linear(300)(x)
x = jax.nn.relu(x)
x = eg.Linear(10)(x)
return x
```
**2.** Create a `Model` from this module and specify additional things like losses, metrics, and optimizers:
```python
import optax optax
import elegy as eg
model = eg.Model(
module=MLP(),
loss=[
eg.losses.Crossentropy(),
eg.regularizers.L2(l=1e-5),
],
metrics=eg.metrics.Accuracy(),
optimizer=optax.rmsprop(1e-3),
)
```
**3.** Train the model using the `fit` method:
```python
model.fit(
inputs=X_train,
labels=y_train,
epochs=100,
steps_per_epoch=200,
batch_size=64,
validation_data=(X_test, y_test),
shuffle=True,
callbacks=[eg.callbacks.TensorBoard("summaries")]
)
```
#### Using Flax
<details>
<summary>Show</summary>
To use Flax with Elegy just create a `flax.linen.Module` and pass it to `Model`.
```python
import jax
import elegy as eg
import optax optax
import flax.linen as nn
class MLP(nn.Module):
@nn.compact
def __call__(self, x, training: bool):
x = nn.Dense(300)(x)
x = jax.nn.relu(x)
x = nn.Dense(10)(x)
return x
model = eg.Model(
modExcerpt of 9,581 characters
Read on GitHubWould you bet a product on this? Bounded 0–100 and slow moving.
matched fp:22b6fdbe7f5f09f6, topic:deep-learning, topic:jax