JVP treatment of mask as a leaf vs pytree

[rdyro]

The mask in jvp might not match the pytree structure of grad_y.

Unfortunately, I don't have an open-source repro at the moment, but these lines assume pytrees of a similar structure, however the mask is a leaf, but the gradients are a pytree

• folx/folx/jvp.py

  Line 183 in d3bf210

  new_masks = broadcast_mask_to_jacobian(out_mask, grad_y)

• folx/folx/jvp.py

  Line 248 in d3bf210

  result_mask = broadcast_mask_to_jacobian(result_mask, grad_y)

Is that correct?

I believe you recently addressed a similar issue here:
5150e33#diff-21e634aa62155f577c8e87e1b851189b4791db79bdb2593cc957ca86e8cde5ccL328

[n-gao]

If you could create a minimal example that would be great!

[rdyro]

I'm struggling to trigger it with a simple case.

In vjp.py get_jacobian_for_reduction returns a single Array mask, but

    grad_tan, out_mask = get_jacobian_for_reduction(laplace_args.jacobian, axes)
    ...
    @functools.partial(jax.vmap, in_axes=0, out_axes=(None, 0))
    def jvp(tangents):
        return jax.jvp(fwd, laplace_args.x, tangents)

    y, y_tangent = jvp(tangent)
    grad_y = tree_take(y_tangent, slice(None, -1), axis=JAC_DIM)
    lapl_y = tree_take(y_tangent, -1, axis=JAC_DIM)
    print(f"{y_tangent=} {grad_y=}")

    new_masks = broadcast_mask_to_jacobian(out_mask, grad_y)

y_tangent comes out as a pytree (in the internal test it's a list[Array] of size 1) and broadcast_mask_to_jacobian expects both arguments to be pytrees of the same structure.

Apply this diff

diff --git a/folx/jvp.py b/folx/jvp.py
index 5d9a73e..0199d6a 100644
--- a/folx/jvp.py
+++ b/folx/jvp.py
@@ -180,8 +180,8 @@ def sparse_jvp(
     grad_y = tree_take(y_tangent, slice(None, -1), axis=JAC_DIM)
     lapl_y = tree_take(y_tangent, -1, axis=JAC_DIM)
 
-    new_masks = broadcast_mask_to_jacobian(out_mask, grad_y)
-    assert grad_y.shape == new_masks.shape, f'{grad_y.shape} != {new_masks.shape}'
+    new_masks = jax.tree.map(functools.partial(broadcast_mask_to_jacobian, out_mask), grad_y)
+    assert jax.tree.all(jax.tree.map(lambda x, y: x.shape == y.shape, grad_y, new_masks))
 
     grad_y = jtu.tree_map(FwdJacobian, grad_y, new_masks)
     return y, grad_y, lapl_y
@@ -245,10 +245,10 @@ def sparse_diag_jvp(
 
     # We need to broadcast the output mask to the shape of the gradient in case the operation
     # included some broadcasting, e.g., (10, 1) * (5,) -> (10, 5)
-    result_mask = broadcast_mask_to_jacobian(result_mask, grad_y)
-    assert grad_y.shape == result_mask.shape, f'{grad_y.shape} != {result_mask.shape}'
+    result_masks = jax.tree.map(functools.partial(broadcast_mask_to_jacobian, result_mask), grad_y)
+    assert jax.tree.all(jax.tree.map(lambda x, y: x.shape == y.shape, grad_y, result_masks))
 
-    grad_y = jtu.tree_map(FwdJacobian, grad_y, result_mask)
+    grad_y = jtu.tree_map(FwdJacobian, grad_y, result_masks)
     return y, grad_y, lapl_y

resolves the internal test case (sorry, I'm working on trying to get an actual repro).

Can you suggest how to trigger this code path? I'm struggling to construct a forward function with all weak jacobians.