jax-optix

Name	jax-optix JSON
Version	0.1.8 JSON
	download
home_page	None
Summary	Zero-overhead functional lensing for JAX PyTrees
upload_time	2025-02-13 16:06:02
maintainer	None
docs_url	None
author	None
requires_python	>=3.12
license	MIT
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 jax-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/49/6d/898a30442b8ab13d2a9ee6c7aed52abb51c86bdf55e7f16b4c42048665ca/jax_optix-0.1.8.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 jax-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": "MIT",
    "summary": "Zero-overhead functional lensing for JAX PyTrees",
    "version": "0.1.8",
    "project_urls": {
        "Homepage": "https://github.com/jonkhler/optix"
    },
    "split_keywords": [
        "equinox",
        " functional",
        " jax",
        " lens",
        " optics"
    ],
    "urls": [
        {
            "comment_text": null,
            "digests": {
                "blake2b_256": "ca6a9590cd31ef28ff7eae7228613e2489774e91dff889c559e97cf851e9bf58",
                "md5": "82271c6bcd8103ed03c2ec44ef98b246",
                "sha256": "3a2a191479a3114a1995252ffc97c4ac16d0b6e7cd0b6ee9457687015321b3a6"
            },
            "downloads": -1,
            "filename": "jax_optix-0.1.8-py3-none-any.whl",
            "has_sig": false,
            "md5_digest": "82271c6bcd8103ed03c2ec44ef98b246",
            "packagetype": "bdist_wheel",
            "python_version": "py3",
            "requires_python": ">=3.12",
            "size": 3692,
            "upload_time": "2025-02-13T16:06:01",
            "upload_time_iso_8601": "2025-02-13T16:06:01.482588Z",
            "url": "https://files.pythonhosted.org/packages/ca/6a/9590cd31ef28ff7eae7228613e2489774e91dff889c559e97cf851e9bf58/jax_optix-0.1.8-py3-none-any.whl",
            "yanked": false,
            "yanked_reason": null
        },
        {
            "comment_text": null,
            "digests": {
                "blake2b_256": "496d898a30442b8ab13d2a9ee6c7aed52abb51c86bdf55e7f16b4c42048665ca",
                "md5": "af3076d42d906a6716166e4a24d886d8",
                "sha256": "48e55e0faa986b934365331c97b1039321940ae339e1dda9126f2c3f2fef2a7a"
            },
            "downloads": -1,
            "filename": "jax_optix-0.1.8.tar.gz",
            "has_sig": false,
            "md5_digest": "af3076d42d906a6716166e4a24d886d8",
            "packagetype": "sdist",
            "python_version": "source",
            "requires_python": ">=3.12",
            "size": 3509,
            "upload_time": "2025-02-13T16:06:02",
            "upload_time_iso_8601": "2025-02-13T16:06:02.758860Z",
            "url": "https://files.pythonhosted.org/packages/49/6d/898a30442b8ab13d2a9ee6c7aed52abb51c86bdf55e7f16b4c42048665ca/jax_optix-0.1.8.tar.gz",
            "yanked": false,
            "yanked_reason": null
        }
    ],
    "upload_time": "2025-02-13 16:06:02",
    "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"
}