Fix sparse Jacobians with empty dependencies and singleton batch axes

folx/api.py

@@ -187,7 +187,7 @@ def dense_array(self) -> Array:
     @property
     def max_n(self) -> int:
         if self.x0_idx is not None:
-            return int(np.max(self.x0_idx))
+            return int(np.max(self.x0_idx, initial=-1))
         return self.data.shape[JAC_DIM]
 
     @property

folx/hessian.py

@@ -371,6 +371,10 @@ def remove_zero_entries(
     if dense_out:
         return lapl_args, out_idx, None
 
+    # find_out_idx squeezes size-one position axes. Restore them before
+    # indexing the original operands, whose last axis is contracted.
+    frame = np.broadcast_shapes(*(arg.x.shape[:-1] for arg in lapl_args.arrays))
+    out_idx = out_idx.reshape(out_idx.shape[JAC_DIM], *frame)
     mask = (out_idx != -1).any(0)
     if mask.sum() > 0.5 * mask.size:
         # this is a heuristic to avoid having unnecessary indexing overhead for

test/test_index_mask.py

@@ -105,9 +105,10 @@ def test_dense_array(shape, k):
     np.testing.assert_allclose(dense, expected, rtol=1e-12, atol=1e-12)
 
 
-def test_dense_array_all_invalid():
+@pytest.mark.parametrize('rows', [0, 2])
+def test_dense_array_all_invalid(rows):
     """A Jacobian without any dependency densifies to a single zero row."""
-    mask = np.full((2, 3), -1)
-    dense = FwdJacobian(jnp.ones((2, 3)), mask).dense_array
+    mask = np.full((rows, 3), -1, dtype=int)
+    dense = FwdJacobian(jnp.ones((rows, 3)), mask).dense_array
     assert dense.shape == (1, 3)
     np.testing.assert_allclose(dense, 0)

test/test_sparse_attention.py

@@ -0,0 +1,41 @@
+"""Sparse contractions must preserve singleton batch axes during indexing."""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+
+from folx import forward_laplacian
+
+
+@pytest.mark.parametrize('batch_shape', [(), (1,), (1, 1), (2,)])
+def test_sparse_attention_batch_axes(batch_shape):
+    rng = np.random.default_rng(123)
+    x = jnp.asarray(rng.normal(size=(*batch_shape, 6, 3)))
+    weights = jnp.asarray(rng.normal(size=(3, 8)))
+
+    def attention(x):
+        q = jnp.tanh(x @ weights).reshape(6, 2, 4)
+        scores = jnp.einsum('ihd,jhd->hij', q, q)
+        return (jax.nn.softmax(scores) ** 2).sum()
+
+    def f(x):
+        if not batch_shape:
+            return attention(x)
+        return jax.vmap(attention)(x.reshape(-1, 6, 3)).sum()
+
+    actual = jax.jit(forward_laplacian(f, 0.5))(x)
+    expected_laplacian = jax.hessian(f)(x).reshape(x.size, x.size).trace()
+    np.testing.assert_allclose(actual.x, f(x), rtol=1e-6, atol=1e-6)
+    np.testing.assert_allclose(
+        actual.jacobian.dense_array.reshape(x.shape),
+        jax.grad(f)(x),
+        rtol=1e-6,
+        atol=1e-6,
+    )
+    np.testing.assert_allclose(
+        actual.laplacian,
+        expected_laplacian,
+        rtol=1e-6,
+        atol=1e-6,
+    )