{"id":17921270,"url":"https://github.com/francois-rozet/inox","last_synced_at":"2025-09-16T09:31:41.844Z","repository":{"id":169817802,"uuid":"645838495","full_name":"francois-rozet/inox","owner":"francois-rozet","description":"Stainless neural networks in JAX","archived":false,"fork":false,"pushed_at":"2024-10-15T14:20:30.000Z","size":161,"stargazers_count":32,"open_issues_count":0,"forks_count":0,"subscribers_count":4,"default_branch":"master","last_synced_at":"2024-12-30T21:41:41.963Z","etag":null,"topics":["deep-learning","jax","neural-networks"],"latest_commit_sha":null,"homepage":"https://inox.readthedocs.io","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/francois-rozet.png","metadata":{"files":{"readme":"README.md","changelog":null,"contributing":"CONTRIBUTING.md","funding":null,"license":"LICENSE","code_of_conduct":null,"threat_model":null,"audit":null,"citation":null,"codeowners":null,"security":null,"support":null,"governance":null,"roadmap":null,"authors":null,"dei":null,"publiccode":null,"codemeta":null}},"created_at":"2023-05-26T14:52:43.000Z","updated_at":"2024-12-18T01:31:24.000Z","dependencies_parsed_at":"2024-01-17T14:47:44.889Z","dependency_job_id":"879690d5-827d-4874-9b90-7825fc165137","html_url":"https://github.com/francois-rozet/inox","commit_stats":null,"previous_names":["francois-rozet/inox"],"tags_count":16,"template":false,"template_full_name":null,"repository_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/francois-rozet%2Finox","tags_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/francois-rozet%2Finox/tags","releases_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/francois-rozet%2Finox/releases","manifests_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/repositories/francois-rozet%2Finox/manifests","owner_url":"https://repos.ecosyste.ms/api/v1/hosts/GitHub/owners/francois-rozet","download_url":"https://codeload.github.com/francois-rozet/inox/tar.gz/refs/heads/master","host":{"name":"GitHub","url":"https://github.com","kind":"github","repositories_count":233243701,"owners_count":18646934,"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","neural-networks"],"created_at":"2024-10-28T20:32:36.590Z","updated_at":"2025-09-16T09:31:36.477Z","avatar_url":"https://github.com/francois-rozet.png","language":"Python","funding_links":[],"categories":[],"sub_categories":[],"readme":"![Inox's banner](https://raw.githubusercontent.com/francois-rozet/inox/master/docs/images/banner.svg)\n\n# Stainless neural networks in JAX\n\nInox is a minimal [JAX](https://github.com/google/jax) library for neural networks with an intuitive [PyTorch](https://github.com/pytorch/pytorch)-like syntax. As with [Equinox](https://github.com/patrick-kidger/equinox), modules are represented as PyTrees, which enables complex architectures, easy manipulations, and functional transformations.\n\nInox aims to be a leaner version of Equinox by only retaining its core features: PyTrees and lifted transformations. In addition, Inox takes inspiration from other projects like [NNX](https://github.com/cgarciae/nnx) and [Serket](https://github.com/ASEM000/serket) to provide a versatile interface. Despite the differences, Inox remains compatible with the Equinox ecosystem, and its components (modules, transformations, ...) are for the most part interchangeable with those of Equinox.\n\n\u003e Inox means \"stainless steel\" in French 🔪\n\n## Installation\n\nThe `inox` package is available on [PyPI](https://pypi.org/project/inox), which means it is installable via `pip`.\n\n```\npip install inox\n```\n\nAlternatively, if you need the latest features, you can install it from the repository.\n\n```\npip install git+https://github.com/francois-rozet/inox\n```\n\n## Getting started\n\nModules are defined with an intuitive PyTorch-like syntax,\n\n```python\nimport jax\nimport inox.nn as nn\n\ninit_key, data_key = jax.random.split(jax.random.key(0))\n\nclass MLP(nn.Module):\n    def __init__(self, key):\n        keys = jax.random.split(key, 3)\n\n        self.l1 = nn.Linear(3, 64, key=keys[0])\n        self.l2 = nn.Linear(64, 64, key=keys[1])\n        self.l3 = nn.Linear(64, 3, key=keys[2])\n        self.relu = nn.ReLU()\n\n    def __call__(self, x):\n        x = self.l1(x)\n        x = self.l2(self.relu(x))\n        x = self.l3(self.relu(x))\n\n        return x\n\nmodel = MLP(init_key)\n```\n\nand are compatible with JAX transformations.\n\n```python\nX = jax.random.normal(data_key, (1024, 3))\nY = jax.numpy.sort(X, axis=-1)\n\n@jax.jit\ndef loss_fn(model, x, y):\n    pred = jax.vmap(model)(x)\n    return jax.numpy.mean((y - pred) ** 2)\n\ngrads = jax.grad(loss_fn)(model, X, Y)\n```\n\nHowever, if a tree contains strings or boolean flags, it becomes incompatible with JAX transformations. For this reason, Inox provides lifted transformations that consider all non-array leaves as static.\n\n```python\nmodel.name = 'stainless'  # not an array\n\n@inox.jit\ndef loss_fn(model, x, y):\n    pred = jax.vmap(model)(x)\n    return jax.numpy.mean((y - pred) ** 2)\n\ngrads = inox.grad(loss_fn)(model, X, Y)\n```\n\nInox also provides a partition mechanism to split the static definition of a module (structure, strings, flags, ...) from its dynamic content (parameters, indices, statistics, ...), which is convenient for updating parameters.\n\n```python\nmodel.mask = jax.numpy.array([1, 0, 1])  # not a parameter\n\nstatic, params, others = model.partition(nn.Parameter)\n\n@jax.jit\ndef loss_fn(params, others, x, y):\n    model = static(arrays, others)\n    pred = jax.vmap(model)(x)\n    return jax.numpy.mean((y - pred) ** 2)\n\ngrads = jax.grad(loss_fn)(params, others, X, Y)\nparams = jax.tree_util.tree_map(lambda p, g: p - 0.01 * g, params, grads)\n\nmodel = static(params, others)\n```\n\nFor more information, check out the documentation and tutorials at [inox.readthedocs.io](https://inox.readthedocs.io).\n\n## Contributing\n\nIf you have a question, an issue or would like to contribute, please read our [contributing guidelines](https://github.com/francois-rozet/inox/blob/master/CONTRIBUTING.md).\n","project_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffrancois-rozet%2Finox","html_url":"https://awesome.ecosyste.ms/projects/github.com%2Ffrancois-rozet%2Finox","lists_url":"https://awesome.ecosyste.ms/api/v1/projects/github.com%2Ffrancois-rozet%2Finox/lists"}