Bump version and add tests for kwargs

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "v0.1.3"
+current_version = "v0.2.0"
 commit = true
 commit_args = "--no-verify"
 tag = true

README.md

@@ -1,5 +1,5 @@
 # sparsejac: Efficient sparse Jacobians using Jax
-`v0.1.3`
+`v0.2.0`
 
 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.
 

pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 
 name = "sparsejac"
-version = "v0.1.3"
+version = "v0.2.0"
 description = "Efficient forward- and reverse-mode sparse Jacobians using Jax."
 keywords = ["jax", "jacobian", "sparse"]
 readme = "README.md"

src/sparsejac/__init__.py

@@ -1,6 +1,6 @@
 """sparsejac - Efficient forward- and reverse-mode sparse Jacobians using Jax."""
 
-__version__ = "v0.1.3"
+__version__ = "v0.2.0"
 __author__ = "Martin Schubert <[email protected]>"
 
 from sparsejac.sparsejac import jacfwd as jacfwd

tests/test_sparsejac.py

@@ -181,6 +181,19 @@ def fn(x):
         onp.testing.assert_array_equal(expected_jac, result_jac.todense())
         onp.testing.assert_array_equal(expected_aux, result_aux)
 
+    def test_kwargs(self):
+        def fn(x, y):
+            return y * jnp.convolve(x, jnp.asarray([1.0, -2.0, 1.0]), mode="same") ** 2
+
+        x = jax.random.uniform(jax.random.PRNGKey(0), shape=(_SIZE,))
+        i, j = jnp.meshgrid(jnp.arange(_SIZE), jnp.arange(_SIZE), indexing="ij")
+        sparsity = (i == j) | ((i - 1) == j) | ((i + 1) == j)
+        sparsity = jsparse.BCOO.fromdense(sparsity)
+
+        result_jac = sparsejac.jacrev(fn, sparsity)(x, y=1)
+        expected_jac = jax.jacrev(fn)(x, y=1)
+        onp.testing.assert_array_equal(expected_jac, result_jac.todense())
+
 
 class JacfwdTest(unittest.TestCase):
     def test_sparsity_shape_validation(self):
@@ -350,6 +363,19 @@ def fn(x):
         onp.testing.assert_array_equal(expected_jac, result_jac.todense())
         onp.testing.assert_array_equal(expected_aux, result_aux)
 
+    def test_kwargs(self):
+        def fn(x, y):
+            return y * jnp.convolve(x, jnp.asarray([1.0, -2.0, 1.0]), mode="same") ** 2
+
+        x = jax.random.uniform(jax.random.PRNGKey(0), shape=(_SIZE,))
+        i, j = jnp.meshgrid(jnp.arange(_SIZE), jnp.arange(_SIZE), indexing="ij")
+        sparsity = (i == j) | ((i - 1) == j) | ((i + 1) == j)
+        sparsity = jsparse.BCOO.fromdense(sparsity)
+
+        result_jac = sparsejac.jacfwd(fn, sparsity)(x, y=1)
+        expected_jac = jax.jacfwd(fn)(x, y=1)
+        onp.testing.assert_array_equal(expected_jac, result_jac.todense())
+
 
 class ConnectivityFromSparsityTest(unittest.TestCase):
     def test_output_connectivity_matches_expected(self):