jax-optix

Name	jax-optix JSON
Version	0.1.4 JSON
	download
home_page	None
Summary	Zero-overhead functional lensing for JAX PyTrees
upload_time	2024-10-28 15:52:24
maintainer	None
docs_url	None
author	None
requires_python	>=3.12
license	None
keywords	equinox functional jax lens optics
VCS
bugtrack_url
requirements	No requirements were recorded.
Travis-CI	No Travis.
coveralls test coverage	No coveralls.

            # Optix 🔍

A functional lensing library for JAX/Equinox, providing a way to focus on and modify nested values within PyTree structures. Optix generates the same HLO code as direct access, ensuring zero overhead.

## Features

- Type-safe lenses for any JAX PyTree structure
- Zero runtime overhead (generates identical HLO code)
- Intuitive API for accessing and modifying nested values
- Complete static typing support

## Example

```python
from optix import focus
import jax.numpy as jnp

# Create a nested PyTree structure
data = MyStruct(
    x=jnp.array([1.0, 2.0]),
    nested=NestedStruct(y=jnp.array(3.0))
)

# Focus on and modify a nested value
result = focus(data).at(lambda x: x.nested.y).apply(jnp.square)
>>> MyStruct(
>>>     x=Array([1., 2.], dtype=float32),
>>>     nested=NestedStruct(
>>>         y=Array(9., dtype=float32)
>>>     )
>>> )
```

## Installation

```bash
pip install optix
```

## License

MIT License

## Credits

Special thanks to [Patrick Kidger](https://kidger.site/) for providing helpful hints and the [Equinox](https://github.com/patrick-kidger/equinox) library, which this project builds upon.

            

Raw data

            {
    "_id": null,
    "home_page": null,
    "name": "jax-optix",
    "maintainer": null,
    "docs_url": null,
    "requires_python": ">=3.12",
    "maintainer_email": null,
    "keywords": "equinox, functional, jax, lens, optics",
    "author": null,
    "author_email": "Jonas K\u00f6hler <jonas.koehler.ks@gmail.com>",
    "download_url": "https://files.pythonhosted.org/packages/4e/c7/743150d5d88e54c1b5d68074d56d8be432ea0a2190693c54b09388014df3/jax_optix-0.1.4.tar.gz",
    "platform": null,
    "description": "# Optix \ud83d\udd0d\n\nA functional lensing library for JAX/Equinox, providing a way to focus on and modify nested values within PyTree structures. Optix generates the same HLO code as direct access, ensuring zero overhead.\n\n## Features\n\n- Type-safe lenses for any JAX PyTree structure\n- Zero runtime overhead (generates identical HLO code)\n- Intuitive API for accessing and modifying nested values\n- Complete static typing support\n\n## Example\n\n```python\nfrom optix import focus\nimport jax.numpy as jnp\n\n# Create a nested PyTree structure\ndata = MyStruct(\n    x=jnp.array([1.0, 2.0]),\n    nested=NestedStruct(y=jnp.array(3.0))\n)\n\n# Focus on and modify a nested value\nresult = focus(data).at(lambda x: x.nested.y).apply(jnp.square)\n>>> MyStruct(\n>>>     x=Array([1., 2.], dtype=float32),\n>>>     nested=NestedStruct(\n>>>         y=Array(9., dtype=float32)\n>>>     )\n>>> )\n```\n\n## Installation\n\n```bash\npip install optix\n```\n\n## License\n\nMIT License\n\n## Credits\n\nSpecial thanks to [Patrick Kidger](https://kidger.site/) for providing helpful hints and the [Equinox](https://github.com/patrick-kidger/equinox) library, which this project builds upon.\n",
    "bugtrack_url": null,
    "license": null,
    "summary": "Zero-overhead functional lensing for JAX PyTrees",
    "version": "0.1.4",
    "project_urls": {
        "Homepage": "https://github.com/jonkhler/optix"
    },
    "split_keywords": [
        "equinox",
        " functional",
        " jax",
        " lens",
        " optics"
    ],
    "urls": [
        {
            "comment_text": "",
            "digests": {
                "blake2b_256": "b32398eb5628a32f7a63e4054be917ba55e590b8c7d313c853ad158e8af1f7d1",
                "md5": "fd5328ee1e43b24b0ae695327de9fbd3",
                "sha256": "6576a4f0da1b631892b39053ff72a5ed9ef62f5cf5e1178530c1acb731fc11a3"
            },
            "downloads": -1,
            "filename": "jax_optix-0.1.4-py3-none-any.whl",
            "has_sig": false,
            "md5_digest": "fd5328ee1e43b24b0ae695327de9fbd3",
            "packagetype": "bdist_wheel",
            "python_version": "py3",
            "requires_python": ">=3.12",
            "size": 3473,
            "upload_time": "2024-10-28T15:52:22",
            "upload_time_iso_8601": "2024-10-28T15:52:22.254370Z",
            "url": "https://files.pythonhosted.org/packages/b3/23/98eb5628a32f7a63e4054be917ba55e590b8c7d313c853ad158e8af1f7d1/jax_optix-0.1.4-py3-none-any.whl",
            "yanked": false,
            "yanked_reason": null
        },
        {
            "comment_text": "",
            "digests": {
                "blake2b_256": "4ec7743150d5d88e54c1b5d68074d56d8be432ea0a2190693c54b09388014df3",
                "md5": "81e6c0f781d1313284940f1b53e1dcd2",
                "sha256": "d35969f3ef67c37bcbf3394e504d94c92499ba214a6db0f13f21190630526052"
            },
            "downloads": -1,
            "filename": "jax_optix-0.1.4.tar.gz",
            "has_sig": false,
            "md5_digest": "81e6c0f781d1313284940f1b53e1dcd2",
            "packagetype": "sdist",
            "python_version": "source",
            "requires_python": ">=3.12",
            "size": 3270,
            "upload_time": "2024-10-28T15:52:24",
            "upload_time_iso_8601": "2024-10-28T15:52:24.388313Z",
            "url": "https://files.pythonhosted.org/packages/4e/c7/743150d5d88e54c1b5d68074d56d8be432ea0a2190693c54b09388014df3/jax_optix-0.1.4.tar.gz",
            "yanked": false,
            "yanked_reason": null
        }
    ],
    "upload_time": "2024-10-28 15:52:24",
    "github": true,
    "gitlab": false,
    "bitbucket": false,
    "codeberg": false,
    "github_user": "jonkhler",
    "github_project": "optix",
    "travis_ci": false,
    "coveralls": false,
    "github_actions": false,
    "lcname": "jax-optix"
}