{"id":13806474,"url":"https://poets-ai.github.io/elegy/","last_synced_at":"2025-05-13T21:33:14.140Z","repository":{"id":38370743,"uuid":"276115232","full_name":"poets-ai/elegy","owner":"poets-ai","description":"A High Level API for Deep Learning in JAX","archived":false,"fork":false,"pushed_at":"2022-12-15T19:23:10.000Z","size":35542,"stargazers_count":469,"open_issues_count":47,"forks_count":32,"subscribers_count":16,"default_branch":"master","last_synced_at":"2024-11-07T11:44:58.409Z","etag":null,"topics":["deep-learning","jax"],"latest_commit_sha":null,"homepage":"https://poets-ai.github.io/elegy/","language":"Python","has_issues":true,"has_wiki":null,"has_pages":null,"mirror_url":null,"source_name":null,"license":"mit","status":null,"scm":"git","pull_requests_enabled":true,"icon_url":"https://github.com/poets-ai.png","metadata":{"files":{"readme":"README.md","changelog":"CHANGELOG.md","contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null}},"created_at":"2020-06-30T14:00:37.000Z","updated_at":"2024-10-12T09:53:02.000Z","dependencies_parsed_at":"2023-01-29T04:16:04.004Z","dependency_job_id":null,"html_url":"https://github.com/poets-ai/elegy","commit_stats":null,"previous_names":[],"tags_count":22,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/poets-ai%2Felegy","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/poets-ai%2Felegy/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/poets-ai%2Felegy/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/poets-ai%2Felegy/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/poets-ai","download_url":"https://codeload.github.com/poets-ai/elegy/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":225260568,"owners_count":17446094,"icon_url":"https://github.com/github.png","version":null,"created_at":"2022-05-30T11:31:42.601Z","updated_at":"2022-07-04T15:15:14.044Z","host_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub","repositories_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories","repository_names_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repository_names","owners_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners"}},"keywords":["deep-learning","jax"],"created_at":"2024-08-04T01:01:12.249Z","updated_at":"2024-11-18T22:32:12.836Z","avatar_url":"https://github.com/poets-ai.png","language":"Python","funding_links":[],"categories":["Libraries"],"sub_categories":[],"readme":"# Elegy\n\n\u003c!-- [![PyPI Status Badge](https://badge.fury.io/py/eg.svg)](https://pypi.org/project/elegy/) --\u003e\n\u003c!-- [![PyPI - Python Version](https://img.shields.io/pypi/pyversions/elegy)](https://pypi.org/project/elegy/) --\u003e\n\u003c!-- [![Documentation](https://img.shields.io/badge/api-reference-blue.svg)](https://poets-ai.github.io/elegy/) --\u003e\n\u003c!-- [![Code style: black](https://img.shields.io/badge/code%20style-black-000000.svg)](https://github.com/psf/black) --\u003e\n[![Coverage](https://img.shields.io/codecov/c/github/poets-ai/elegy?color=%2334D058)](https://codecov.io/gh/poets-ai/elegy)\n[![Status](https://github.com/poets-ai/elegy/workflows/GitHub%20CI/badge.svg)](https://github.com/poets-ai/elegy/actions?query=workflow%3A%22GitHub+CI%22)\n[![Contributions welcome](https://img.shields.io/badge/contributions-welcome-brightgreen.svg?style=flat)](https://github.com/poets-ai/elegy/issues)\n\n______________________________________________________________________\n\n_A High Level API for Deep Learning in JAX_\n\n#### Main Features\n\n- 😀 **Easy-to-use**: Elegy provides a Keras-like high-level API that makes it very easy to use for most common tasks.\n- 💪‍ **Flexible**: Elegy provides a Pytorch Lightning-like low-level API that offers maximum flexibility when needed.\n- 🔌 **Compatible**: Elegy supports various frameworks and data sources including Flax \u0026 Haiku Modules, Optax Optimizers, TensorFlow Datasets, Pytorch DataLoaders, and more.\n\u003c!-- - 🤷 **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. --\u003e\n\nElegy is built on top of [Treex](https://github.com/cgarciae/treex) and [Treeo](https://github.com/cgarciae/treeo) and reexports their APIs for convenience.\n\n\n[Getting Started](https://poets-ai.github.io/elegy/getting-started/high-level-api) | [Examples](/examples) | [Documentation](https://poets-ai.github.io/elegy)\n\n\n## What is included?\n* A `Model` class with an Estimator-like API.\n* A `callbacks` module with common Keras callbacks.\n\n**From Treex**\n\n* A `Module` class.\n* A `nn` module for with common layers.\n* A `losses` module with common loss functions.\n* A `metrics` module with common metrics.\n\n## Installation\n\nInstall using pip:\n\n```bash\npip install elegy\n```\n\nFor 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.\n\n## Quick Start: High-level API\n\nElegy's high-level API provides a straightforward interface you can use by implementing the following steps:\n\n**1.** Define the architecture inside a `Module`:\n\n```python\nimport jax\nimport elegy as eg\n\nclass MLP(eg.Module):\n    @eg.compact\n    def __call__(self, x):\n        x = eg.Linear(300)(x)\n        x = jax.nn.relu(x)\n        x = eg.Linear(10)(x)\n        return x\n```\n\n**2.** Create a `Model` from this module and specify additional things like losses, metrics, and optimizers:\n\n```python\nimport optax optax\nimport elegy as eg\n\nmodel = eg.Model(\n    module=MLP(),\n    loss=[\n        eg.losses.Crossentropy(),\n        eg.regularizers.L2(l=1e-5),\n    ],\n    metrics=eg.metrics.Accuracy(),\n    optimizer=optax.rmsprop(1e-3),\n)\n```\n\n**3.** Train the model using the `fit` method:\n\n```python\nmodel.fit(\n    inputs=X_train,\n    labels=y_train,\n    epochs=100,\n    steps_per_epoch=200,\n    batch_size=64,\n    validation_data=(X_test, y_test),\n    shuffle=True,\n    callbacks=[eg.callbacks.TensorBoard(\"summaries\")]\n)\n```\n#### Using Flax\n\n\u003cdetails\u003e\n\u003csummary\u003eShow\u003c/summary\u003e\n\nTo use Flax with Elegy just create a `flax.linen.Module` and pass it to `Model`.\n\n```python\nimport jax\nimport elegy as eg\nimport optax optax\nimport flax.linen as nn\n\nclass MLP(nn.Module):\n    @nn.compact\n    def __call__(self, x, training: bool):\n        x = nn.Dense(300)(x)\n        x = jax.nn.relu(x)\n        x = nn.Dense(10)(x)\n        return x\n\n\nmodel = eg.Model(\n    module=MLP(),\n    loss=[\n        eg.losses.Crossentropy(),\n        eg.regularizers.L2(l=1e-5),\n    ],\n    metrics=eg.metrics.Accuracy(),\n    optimizer=optax.rmsprop(1e-3),\n)\n```\n\nAs shown here, Flax Modules can optionally request a `training` argument to `__call__` which will be provided by Elegy / Treex. \n\n\u003c/details\u003e\n\n#### Using Haiku\n\n\u003cdetails\u003e\n\u003csummary\u003eShow\u003c/summary\u003e\n\nTo use Haiku with Elegy do the following: \n\n* Create a `forward` function.\n* Create a `TransformedWithState` object by feeding `forward` to `hk.transform_with_state`.\n* Pass your `TransformedWithState`  to `Model`.\n\nYou can also optionally create your own `hk.Module` and use it in `forward` if needed. Putting everything together should look like this:\n\n```python\nimport jax\nimport elegy as eg\nimport optax optax\nimport haiku as hk\n\n\ndef forward(x, training: bool):\n    x = hk.Linear(300)(x)\n    x = jax.nn.relu(x)\n    x = hk.Linear(10)(x)\n    return x\n\n\nmodel = eg.Model(\n    module=hk.transform_with_state(forward),\n    loss=[\n        eg.losses.Crossentropy(),\n        eg.regularizers.L2(l=1e-5),\n    ],\n    metrics=eg.metrics.Accuracy(),\n    optimizer=optax.rmsprop(1e-3),\n)\n```\n\nAs shown here, `forward` can optionally request a `training` argument which will be provided by Elegy / Treex. \n\n\u003c/details\u003e\n\n## Quick Start: Low-level API\n\nElegy's low-level API lets you explicitly define what goes on during training, testing, and inference. Let's define our own custom `Model` to implement a `LinearClassifier` with pure JAX:\n\n**1.** Define a custom `init_step` method:\n\n```python\nclass LinearClassifier(eg.Model):\n    # use treex's API to declare parameter nodes\n    w: jnp.ndarray = eg.Parameter.node()\n    b: jnp.ndarray = eg.Parameter.node()\n\n    def init_step(self, key: jnp.ndarray, inputs: jnp.ndarray):\n        self.w = jax.random.uniform(\n            key=key,\n            shape=[features_in, 10],\n        )\n        self.b = jnp.zeros([10])\n\n        self.optimizer = self.optimizer.init(self)\n\n        return self\n```\nHere we declared the parameters `w` and `b` using Treex's `Parameter.node()` for pedagogical reasons, however normally you don't have to do this since you typically use a sub-`Module` instead.\n\n**2.** Define a custom `test_step` method:\n```python\n    def test_step(self, inputs, labels):\n        # flatten + scale\n        inputs = jnp.reshape(inputs, (inputs.shape[0], -1)) / 255\n\n        # forward\n        logits = jnp.dot(inputs, self.w) + self.b\n\n        # crossentropy loss\n        target = jax.nn.one_hot(labels[\"target\"], 10)\n        loss = optax.softmax_cross_entropy(logits, target).mean()\n\n        # metrics\n        logs = dict(\n            acc=jnp.mean(jnp.argmax(logits, axis=-1) == labels[\"target\"]),\n            loss=loss,\n        )\n\n        return loss, logs, self\n```\n\n**3.** Instantiate our `LinearClassifier` with an optimizer:\n\n```python\nmodel = LinearClassifier(\n    optimizer=optax.rmsprop(1e-3),\n)\n```\n\n**4.** Train the model using the `fit` method:\n\n```python\nmodel.fit(\n    inputs=X_train,\n    labels=y_train,\n    epochs=100,\n    steps_per_epoch=200,\n    batch_size=64,\n    validation_data=(X_test, y_test),\n    shuffle=True,\n    callbacks=[eg.callbacks.TensorBoard(\"summaries\")]\n)\n```\n\n#### Using other JAX Frameworks\n\n\u003cdetails\u003e\n\u003csummary\u003eShow\u003c/summary\u003e\n\nIt is straightforward to integrate other functional JAX libraries with this\nlow-level API, here is an example with Flax:\n\n```python\nimport elegy as eg\nimport flax.linen as nn\n\nclass LinearClassifier(eg.Model):\n    params: Mapping[str, Any] = eg.Parameter.node()\n    batch_stats: Mapping[str, Any] = eg.BatchStat.node()\n    next_key: eg.KeySeq\n\n    def __init__(self, module: nn.Module, **kwargs):\n        self.flax_module = module\n        super().__init__(**kwargs)\n\n    def init_step(self, key, inputs):\n        self.next_key = eg.KeySeq(key)\n\n        variables = self.flax_module.init(\n            {\"params\": self.next_key(), \"dropout\": self.next_key()}, x\n        )\n        self.params = variables[\"params\"]\n        self.batch_stats = variables[\"batch_stats\"]\n\n        self.optimizer = self.optimizer.init(self.parameters())\n\n    def test_step(self, inputs, labels):\n        # forward\n        variables = dict(\n            params=self.params,\n            batch_stats=self.batch_stats,\n        )\n        logits, variables = self.flax_module.apply(\n            variables,\n            inputs, \n            rngs={\"dropout\": self.next_key()}, \n            mutable=True,\n        )\n        self.batch_stats = variables[\"batch_stats\"]\n        \n        # loss\n        target = jax.nn.one_hot(labels[\"target\"], 10)\n        loss = optax.softmax_cross_entropy(logits, target).mean()\n\n        # logs\n        logs = dict(\n            accuracy=accuracy,\n            loss=loss,\n        )\n        return loss, logs, self\n```\n\n\u003c/details\u003e\n\n### Examples\n\nCheck out the [/example](/examples) directory for some inspiration. To run an example, first install some requirements:\n\n```bash\npip install -r examples/requirements.txt\n```\n\nAnd the run it normally with python e.g.\n\n```bash\npython examples/flax/mnist_vae.py\n```\n\n## Contributing\n\nIf your are interested in helping improve Elegy check out the [Contributing Guide](https://poets-ai.github.io/elegy/guides/contributing).\n\n## Sponsors 💚\n* [Quansight](https://www.quansight.com) - paid development time\n\n## Citing Elegy\n\n\n**BibTeX**\n\n```\n@software{elegy2020repository,\n\ttitle        = {Elegy: A High Level API for Deep Learning in JAX},\n\tauthor       = {PoetsAI},\n\tyear         = 2021,\n\turl          = {https://github.com/poets-ai/elegy},\n\tversion      = {0.8.1}\n}\n```\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/poets-ai.github.io%2Felegy%2F","html_url":"https://awesome.ecosyste.ms/projects/poets-ai.github.io%2Felegy%2F","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/poets-ai.github.io%2Felegy%2F/lists"}