Name | paramax JSON |
Version |
0.0.0
JSON |
| download |
home_page | None |
Summary | Parameterizations and parameter constraints for JAX PyTrees. |
upload_time | 2024-10-21 12:19:07 |
maintainer | None |
docs_url | None |
author | None |
requires_python | >=3.10 |
license | The MIT License (MIT) Copyright (c) 2022 Daniel Ward Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. |
keywords |
equinox
jax
neural-networks
|
VCS |
|
bugtrack_url |
|
requirements |
No requirements were recorded.
|
Travis-CI |
No Travis.
|
coveralls test coverage |
No coveralls.
|
Paramax
============
Parameterizations and constraints for JAX PyTrees
-----------------------------------------------------------------------
Paramax allows applying custom constraints or behaviors to PyTree components,
using unwrappable placeholders. This can be used for
- Enforcing positivity (e.g., scale parameters)
- Structured matrices (triangular, symmetric, etc.)
- Applying tricks like weight normalization
- Marking components as non-trainable
Some benefits of the unwrappable pattern:
- It allows parameterizations to be computed once for a model (e.g. at the top of the
loss function).
- It is flexible, e.g. allowing custom parameterizations to be applied to PyTrees
from external libraries
- It is concise
If you found the package useful, please consider giving it a star on github, and if you
create ``AbstractUnwrappable``s that may be of interest to others, a pull request would
be much appreciated!
## Documentation
Documentation available [here](https://danielward27.github.io/paramax/).
## Installation
```bash
pip install paramax
```
## Example
```python
>>> import paramax
>>> import jax.numpy as jnp
>>> scale = paramax.Parameterize(jnp.exp, jnp.log(jnp.ones(3))) # Enforce positivity
>>> paramax.unwrap(("abc", 1, scale))
('abc', 1, Array([1., 1., 1.], dtype=float32))
```
## Alternative parameterization patterns
Using properties to access parameterized model components is common but has drawbacks:
- Parameterizations are tied to class definition, limiting flexibility e.g. this
cannot be used on PyTrees from external libraries
- It can become verbose with many parameters
- It often leads to repeatedly computing the parameterization
## Related
- We make use of the [Equinox](https://arxiv.org/abs/2111.00254) package, to register
the PyTrees used in the package
- This package spawned out of a need for a simple method to apply parameter constraints
in the distributions package [flowjax](https://github.com/danielward27/flowjax)
Raw data
{
"_id": null,
"home_page": null,
"name": "paramax",
"maintainer": null,
"docs_url": null,
"requires_python": ">=3.10",
"maintainer_email": null,
"keywords": "equinox, jax, neural-networks",
"author": null,
"author_email": "Daniel Ward <danielward27@outlook.com>",
"download_url": "https://files.pythonhosted.org/packages/db/9e/a8a69135a884ed48f9dffbf0ed3bb5cdddcab74607d239321ebe2681ed2e/paramax-0.0.0.tar.gz",
"platform": null,
"description": "\nParamax\n============\nParameterizations and constraints for JAX PyTrees\n-----------------------------------------------------------------------\n\nParamax allows applying custom constraints or behaviors to PyTree components,\nusing unwrappable placeholders. This can be used for\n- Enforcing positivity (e.g., scale parameters)\n- Structured matrices (triangular, symmetric, etc.)\n- Applying tricks like weight normalization\n- Marking components as non-trainable\n\nSome benefits of the unwrappable pattern:\n- It allows parameterizations to be computed once for a model (e.g. at the top of the\n loss function).\n- It is flexible, e.g. allowing custom parameterizations to be applied to PyTrees\n from external libraries\n- It is concise\n\nIf you found the package useful, please consider giving it a star on github, and if you\ncreate ``AbstractUnwrappable``s that may be of interest to others, a pull request would\nbe much appreciated!\n\n## Documentation\n\nDocumentation available [here](https://danielward27.github.io/paramax/).\n\n## Installation\n```bash\npip install paramax\n```\n\n## Example\n```python\n>>> import paramax\n>>> import jax.numpy as jnp\n>>> scale = paramax.Parameterize(jnp.exp, jnp.log(jnp.ones(3))) # Enforce positivity\n>>> paramax.unwrap((\"abc\", 1, scale))\n('abc', 1, Array([1., 1., 1.], dtype=float32))\n```\n\n## Alternative parameterization patterns\nUsing properties to access parameterized model components is common but has drawbacks:\n- Parameterizations are tied to class definition, limiting flexibility e.g. this\n cannot be used on PyTrees from external libraries\n- It can become verbose with many parameters\n- It often leads to repeatedly computing the parameterization\n\n## Related\n- We make use of the [Equinox](https://arxiv.org/abs/2111.00254) package, to register\nthe PyTrees used in the package\n- This package spawned out of a need for a simple method to apply parameter constraints\n in the distributions package [flowjax](https://github.com/danielward27/flowjax)\n",
"bugtrack_url": null,
"license": "The MIT License (MIT) Copyright (c) 2022 Daniel Ward Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the \"Software\"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.",
"summary": "Parameterizations and parameter constraints for JAX PyTrees.",
"version": "0.0.0",
"project_urls": {
"documentation": "https://danielward27.github.io/paramax/index.html",
"repository": "https://github.com/danielward27/paramax"
},
"split_keywords": [
"equinox",
" jax",
" neural-networks"
],
"urls": [
{
"comment_text": "",
"digests": {
"blake2b_256": "20f87b1f5b84f84e43face7349ac2df23f0f75d36b5b4a1c1c0305edce82bcc8",
"md5": "33223f6f3c48797321d3551e6057bedf",
"sha256": "02d0120e626de300680a1661b138feeba14b418031d6e976fd679db5fd03509a"
},
"downloads": -1,
"filename": "paramax-0.0.0-py3-none-any.whl",
"has_sig": false,
"md5_digest": "33223f6f3c48797321d3551e6057bedf",
"packagetype": "bdist_wheel",
"python_version": "py3",
"requires_python": ">=3.10",
"size": 6720,
"upload_time": "2024-10-21T12:19:06",
"upload_time_iso_8601": "2024-10-21T12:19:06.334656Z",
"url": "https://files.pythonhosted.org/packages/20/f8/7b1f5b84f84e43face7349ac2df23f0f75d36b5b4a1c1c0305edce82bcc8/paramax-0.0.0-py3-none-any.whl",
"yanked": false,
"yanked_reason": null
},
{
"comment_text": "",
"digests": {
"blake2b_256": "db9ea8a69135a884ed48f9dffbf0ed3bb5cdddcab74607d239321ebe2681ed2e",
"md5": "12b85e600a920282f9ee084edd025136",
"sha256": "871e8809726f535506fdd9cc14116e88adfc66598fab1e309f990f4c7dcbd2bf"
},
"downloads": -1,
"filename": "paramax-0.0.0.tar.gz",
"has_sig": false,
"md5_digest": "12b85e600a920282f9ee084edd025136",
"packagetype": "sdist",
"python_version": "source",
"requires_python": ">=3.10",
"size": 10081,
"upload_time": "2024-10-21T12:19:07",
"upload_time_iso_8601": "2024-10-21T12:19:07.170816Z",
"url": "https://files.pythonhosted.org/packages/db/9e/a8a69135a884ed48f9dffbf0ed3bb5cdddcab74607d239321ebe2681ed2e/paramax-0.0.0.tar.gz",
"yanked": false,
"yanked_reason": null
}
],
"upload_time": "2024-10-21 12:19:07",
"github": true,
"gitlab": false,
"bitbucket": false,
"codeberg": false,
"github_user": "danielward27",
"github_project": "paramax",
"travis_ci": false,
"coveralls": false,
"github_actions": true,
"lcname": "paramax"
}