Announcing Elegy!

[cgarciae]

Hey, I created this framework called Elegy with a couple of friends, its Keras-like Deep Learning framework based on Jax + Haiku. Its still in very early stage but we would love to get feedback. We've been using it for a few days and we think it greatly simplifies creating DL models with Jax.

https://poets-ai.github.io/elegy/

Feel free to close the issue.

[hawkinsp]

Awesome!

(Rather than close this issue, I converted it to a GitHub discussion where the conversation can continue...)

[castellwhite]

The Next Generation of the ML tools!

[8bitmp3]

This is awesome @cgarciae SavedModels soon?

From the README https://github.com/poets-ai/elegy :

!pip install elegy

import flax.linen as nn
import jax
import elegy, optax

class MLP(nn.Module):
    @nn.compact
    def call(self, x):
        x = nn.Dense(300)(x)
        x = jax.nn.relu(x)
        x = nn.Dense(10)(x)
        return x

model = elegy.Model(
    module=MLP(),
    loss=[
        elegy.losses.SparseCategoricalCrossentropy(from_logits=True),
        elegy.regularizers.GlobalL2(l=1e-5),
    ],
    metrics=elegy.metrics.SparseCategoricalAccuracy(),
    optimizer=optax.rmsprop(1e-3),
)

model.fit(
    x=X_train,
    y=y_train,
    epochs=100,
    steps_per_epoch=200,
    batch_size=64,
    validation_data=(X_test, y_test),
    shuffle=True,
    callbacks=[elegy.callbacks.TensorBoard("summaries")]
)

@cgarciae 👏 This is great work.

WDYT? @hawkinsp @mattjj @avital @marcvanzee @levskaya @dynamicwebpaige and all the teams 👍

High-level Elegy: https://poets-ai.github.io/elegy/getting-started-high-level-api/
Low-level Elegy: https://poets-ai.github.io/elegy/getting-started-low-level-api/