sparsejac

Name	sparsejac JSON
Version	0.1.3 JSON
	download
home_page
Summary	Efficient forward- and reverse-mode sparse Jacobians using Jax.
upload_time	2024-02-18 17:07:45
maintainer
docs_url	None
author
requires_python	>=3.7
license	MIT License Copyright (c) 2023 Martin Schubert 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	jax jacobian sparse
VCS
bugtrack_url
requirements	No requirements were recorded.
Travis-CI	No Travis.
coveralls test coverage	No coveralls.

            # sparsejac: Efficient forward- and reverse-mode sparse Jacobians using Jax.
`v0.1.3`

Sparse Jacobians are frequently encountered in the simulation of physical systems. Jax tranformations `jacfwd` and `jacrev` make it easy to compute dense Jacobians, but these are wasteful when the Jacobian is sparse. `sparsejac` provides a function to more efficiently compute the Jacobian if its sparsity is known. It makes use of the recently-introduced `jax.experimental.sparse` module.

## Install
```
pip install sparsejac
```

## Example
A trivial example with a diagonal Jacobian follows:

```python
fn = lambda x: x**2
x = jax.random.uniform(jax.random.PRNGKey(0), shape=(10000,))

@jax.jit
def sparse_jacrev_fn(x):
  with jax.ensure_compile_time_eval():
    sparsity = jax.experimental.sparse.BCOO.fromdense(jnp.eye(10000))
    jacrev_fn = sparsejac.jacrev(fn, sparsity=sparsity)
  return jacrev_fn(x)

dense_jacrev_fn = jax.jit(jax.jacrev(fn))

assert jnp.all(sparse_jacrev_fn(x).todense() == dense_jacrev_fn(x))

%timeit sparse_jacrev_fn(x).block_until_ready()
%timeit dense_jacrev_fn(x).block_until_ready()
```

And, the performance improvement can easily be seen:

```
93.1 µs ± 17.2 µs per loop (mean ± std. dev. of 7 runs, 10000 loops each)
182 ms ± 26.9 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
```

            

Raw data

            {
    "_id": null,
    "home_page": "",
    "name": "sparsejac",
    "maintainer": "",
    "docs_url": null,
    "requires_python": ">=3.7",
    "maintainer_email": "Martin Schubert <mfschubert@gmail.com>",
    "keywords": "jax,jacobian,sparse",
    "author": "",
    "author_email": "Martin Schubert <mfschubert@gmail.com>",
    "download_url": "https://files.pythonhosted.org/packages/be/2e/2f666e974d3566cd2a6276d465d10f829e3583f927c301cb8e20d5848838/sparsejac-0.1.3.tar.gz",
    "platform": null,
    "description": "# sparsejac: Efficient forward- and reverse-mode sparse Jacobians using Jax.\n`v0.1.3`\n\nSparse Jacobians are frequently encountered in the simulation of physical systems. Jax tranformations `jacfwd` and `jacrev` make it easy to compute dense Jacobians, but these are wasteful when the Jacobian is sparse. `sparsejac` provides a function to more efficiently compute the Jacobian if its sparsity is known. It makes use of the recently-introduced `jax.experimental.sparse` module.\n\n## Install\n```\npip install sparsejac\n```\n\n## Example\nA trivial example with a diagonal Jacobian follows:\n\n```python\nfn = lambda x: x**2\nx = jax.random.uniform(jax.random.PRNGKey(0), shape=(10000,))\n\n@jax.jit\ndef sparse_jacrev_fn(x):\n  with jax.ensure_compile_time_eval():\n    sparsity = jax.experimental.sparse.BCOO.fromdense(jnp.eye(10000))\n    jacrev_fn = sparsejac.jacrev(fn, sparsity=sparsity)\n  return jacrev_fn(x)\n\ndense_jacrev_fn = jax.jit(jax.jacrev(fn))\n\nassert jnp.all(sparse_jacrev_fn(x).todense() == dense_jacrev_fn(x))\n\n%timeit sparse_jacrev_fn(x).block_until_ready()\n%timeit dense_jacrev_fn(x).block_until_ready()\n```\n\nAnd, the performance improvement can easily be seen:\n\n```\n93.1 \u00b5s \u00b1 17.2 \u00b5s per loop (mean \u00b1 std. dev. of 7 runs, 10000 loops each)\n182 ms \u00b1 26.9 ms per loop (mean \u00b1 std. dev. of 7 runs, 10 loops each)\n```\n",
    "bugtrack_url": null,
    "license": "MIT License  Copyright (c) 2023 Martin Schubert  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": "Efficient forward- and reverse-mode sparse Jacobians using Jax.",
    "version": "0.1.3",
    "project_urls": null,
    "split_keywords": [
        "jax",
        "jacobian",
        "sparse"
    ],
    "urls": [
        {
            "comment_text": "",
            "digests": {
                "blake2b_256": "9bdcb838506896332e920764d7898ec0c33879b9eb8ef39f782f3eb656d57dae",
                "md5": "4721dee8231f1876c5837ac4a1730297",
                "sha256": "0016f226ebec7e4173d41c2d8b7d97c6904f5d327b5db53ce3cd72a6d63eba73"
            },
            "downloads": -1,
            "filename": "sparsejac-0.1.3-py3-none-any.whl",
            "has_sig": false,
            "md5_digest": "4721dee8231f1876c5837ac4a1730297",
            "packagetype": "bdist_wheel",
            "python_version": "py3",
            "requires_python": ">=3.7",
            "size": 6482,
            "upload_time": "2024-02-18T17:07:44",
            "upload_time_iso_8601": "2024-02-18T17:07:44.403715Z",
            "url": "https://files.pythonhosted.org/packages/9b/dc/b838506896332e920764d7898ec0c33879b9eb8ef39f782f3eb656d57dae/sparsejac-0.1.3-py3-none-any.whl",
            "yanked": false,
            "yanked_reason": null
        },
        {
            "comment_text": "",
            "digests": {
                "blake2b_256": "be2e2f666e974d3566cd2a6276d465d10f829e3583f927c301cb8e20d5848838",
                "md5": "69ee91e1140e2d9a9a101347b101bca9",
                "sha256": "8aee7fcd1834ebbd1414000111eef9882622ace346719175a02f7f6d4141edeb"
            },
            "downloads": -1,
            "filename": "sparsejac-0.1.3.tar.gz",
            "has_sig": false,
            "md5_digest": "69ee91e1140e2d9a9a101347b101bca9",
            "packagetype": "sdist",
            "python_version": "source",
            "requires_python": ">=3.7",
            "size": 8064,
            "upload_time": "2024-02-18T17:07:45",
            "upload_time_iso_8601": "2024-02-18T17:07:45.543056Z",
            "url": "https://files.pythonhosted.org/packages/be/2e/2f666e974d3566cd2a6276d465d10f829e3583f927c301cb8e20d5848838/sparsejac-0.1.3.tar.gz",
            "yanked": false,
            "yanked_reason": null
        }
    ],
    "upload_time": "2024-02-18 17:07:45",
    "github": false,
    "gitlab": false,
    "bitbucket": false,
    "codeberg": false,
    "lcname": "sparsejac"
}