more robust indexing handling

Author: n-gao

folx/jvp.py

@@ -275,7 +275,10 @@ def sparse_index_jvp(
         # https://github.com/google/jax/pull/3370
         with jax.ensure_compile_time_eval():
             extra_filled = jtu.tree_map(
-                lambda x: jnp.full(x.shape, -1, dtype=jnp.int32), extra_args
+                lambda x: jnp.full(
+                    x.shape if isinstance(x, jax.Array) else (), -1, dtype=jnp.int32
+                ),
+                extra_args,
             )
 
             def _merged_fwd(*args):

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "folx"
-version = "0.2.18"
+version = "0.2.19"
 description = "Forward Laplacian for JAX"
 authors = [
     { name = "Nicholas Gao", email = "n.gao@tum.de" },