Package: elegy
ElegyA High Level API for Deep Learning in JAXMain 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.Elegy is built on top ofTreexandTreeoand reexports their APIs for convenience.Getting Started|Examples|DocumentationWhat is included?AModelclass with an Estimator-like API.Acallbacksmodule with common Keras callbacks.From TreexAModuleclass.Annmodule for with common layers.Alossesmodule with common loss functions.Ametricsmodule with common metrics.InstallationInstall using pip:pipinstallelegyFor Windows users, we recommend the Windows subsystem for Linux 2WSL2sincejaxdoes not support it yet.Quick Start: High-level APIElegy's high-level API provides a straightforward interface you can use by implementing the following steps:1.Define the architecture inside aModule:importjaximportelegyasegclassMLP(eg.Module):@eg.compactdef__call__(self,x):x=eg.Linear(300)(x)x=jax.nn.relu(x)x=eg.Linear(10)(x)returnx2.Create aModelfrom this module and specify additional things like losses, metrics, and optimizers:importoptaxoptaximportelegyasegmodel=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 thefitmethod: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 FlaxShowTo use Flax with Elegy just create aflax.linen.Moduleand pass it toModel.importjaximportelegyasegimportoptaxoptaximportflax.linenasnnclassMLP(nn.Module):@nn.compactdef__call__(self,x,training:bool):x=nn.Dense(300)(x)x=jax.nn.relu(x)x=nn.Dense(10)(x)returnxmodel=eg.Model(module=MLP(),loss=[eg.losses.Crossentropy(),eg.regularizers.L2(l=1e-5),],metrics=eg.metrics.Accuracy(),optimizer=optax.rmsprop(1e-3),)As shown here, Flax Modules can optionally request atrainingargument to__call__which will be provided by Elegy / Treex.Using HaikuShowTo use Haiku with Elegy do the following:Create aforwardfunction.Create aTransformedWithStateobject by feedingforwardtohk.transform_with_state.Pass yourTransformedWithStatetoModel.You can also optionally create your ownhk.Moduleand use it inforwardif needed. Putting everything together should look like this:importjaximportelegyasegimportoptaxoptaximporthaikuashkdefforward(x,training:bool):x=hk.Linear(300)(x)x=jax.nn.relu(x)x=hk.Linear(10)(x)returnxmodel=eg.Model(module=hk.transform_with_state(forward),loss=[eg.losses.Crossentropy(),eg.regularizers.L2(l=1e-5),],metrics=eg.metrics.Accuracy(),optimizer=optax.rmsprop(1e-3),)As shown here,forwardcan optionally request atrainingargument which will be provided by Elegy / Treex.Quick Start: Low-level APIElegy's low-level API lets you explicitly define what goes on during training, testing, and inference. Let's define our own customModelto implement aLinearClassifierwith pure JAX:1.Define a custominit_stepmethod:classLinearClassifier(eg.Model):# use treex's API to declare parameter nodesw:jnp.ndarray=eg.Parameter.node()b:jnp.ndarray=eg.Parameter.node()definit_step(self,key:jnp.ndarray,inputs:jnp.ndarray):self.w=jax.random.uniform(key=key,shape=[features_in,10],)self.b=jnp.zeros([10])self.optimizer=self.optimizer.init(self)returnselfHere we declared the parameterswandbusing Treex'sParameter.node()for pedagogical reasons, however normally you don't have to do this since you typically use a sub-Moduleinstead.2.Define a customtest_stepmethod:deftest_step(self,inputs,labels):# flatten + scaleinputs=jnp.reshape(inputs,(inputs.shape[0],-1))/255# forwardlogits=jnp.dot(inputs,self.w)+self.b# crossentropy losstarget=jax.nn.one_hot(labels["target"],10)loss=optax.softmax_cross_entropy(logits,target).mean()# metricslogs=dict(acc=jnp.mean(jnp.argmax(logits,axis=-1)==labels["target"]),loss=loss,)returnloss,logs,self3.Instantiate ourLinearClassifierwith an optimizer:model=LinearClassifier(optimizer=optax.rmsprop(1e-3),)4.Train the model using thefitmethod: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 other JAX FrameworksShowIt is straightforward to integrate other functional JAX libraries with this
low-level API, here is an example with Flax:importelegyasegimportflax.linenasnnclassLinearClassifier(eg.Model):params:Mapping[str,Any]=eg.Parameter.node()batch_stats:Mapping[str,Any]=eg.BatchStat.node()next_key:eg.KeySeqdef__init__(self,module:nn.Module,**kwargs):self.flax_module=modulesuper().__init__(**kwargs)definit_step(self,key,inputs):self.next_key=eg.KeySeq(key)variables=self.flax_module.init({"params":self.next_key(),"dropout":self.next_key()},x)self.params=variables["params"]self.batch_stats=variables["batch_stats"]self.optimizer=self.optimizer.init(self.parameters())deftest_step(self,inputs,labels):# forwardvariables=dict(params=self.params,batch_stats=self.batch_stats,)logits,variables=self.flax_module.apply(variables,inputs,rngs={"dropout":self.next_key()},mutable=True,)self.batch_stats=variables["batch_stats"]# losstarget=jax.nn.one_hot(labels["target"],10)loss=optax.softmax_cross_entropy(logits,target).mean()# logslogs=dict(accuracy=accuracy,loss=loss,)returnloss,logs,selfExamplesCheck out the/exampledirectory for some inspiration. To run an example, first install some requirements:pipinstall-rexamples/requirements.txtAnd the run it normally with python e.g.pythonexamples/flax/mnist_vae.pyContributingIf your are interested in helping improve Elegy check out theContributing Guide.Sponsors 💚Quansight- paid development timeCiting ElegyBibTeX@software{elegy2020repository,
	title        = {Elegy: A High Level API for Deep Learning in JAX},
	author       = {PoetsAI},
	year         = 2021,
	url          = {https://github.com/poets-ai/elegy},
	version      = {0.8.1}
}