InconclusiveDimensionOperation with complicated vjp

[smartalecH]

I'm trying to run a vjp on a function that concatenates (then flattens) multiple copies of a vector:

from jax import numpy as npj
from jax import vjp, vmap

N0 = 15
N1 = 10

def mapping(x):
    return npj.vstack([x for k in range(N1)]).flatten()

x = npj.ones((N0,))
dJ_du = 0.5*npj.ones((N0*N1,))

y, vjp_fun = vjp(mapping, x)
mg, = vmap(vjp_fun)(dJ_du)
print(mg)

I get an InconclusiveDimensionOperation error with the following (direct) trace:

Traceback (most recent call last):
  File "/Users/alechammond/Projects/grating_study/temp.py", line 14, in <module>
    mg, = vmap(vjp_fun)(dJ_du)
  File "/Users/alechammond/GitHub/jax/jax/core.py", line 1274, in divide_shape_sizes
    raise InconclusiveDimensionOperation(f"Cannot divide evenly the sizes of shapes {tuple(s1)} and {tuple(s2)}")
jax.core.InconclusiveDimensionOperation: Cannot divide evenly the sizes of shapes (150,) and (150, 10, 15)

Note we can do s2//s1 but not s1//s2 (what's needed here). The issue is actually with the flatten step. If I keep the concatenated array unraveled (and manually reshape the vector in the vjp) then it backpropagates just fine.

It seems like something is off with the dimension trace (I admit, the forward function is a bit wonky). Is this expected (or is there a way around this)?

I realize my MWE seems rather trivial and there are much better ways to do the same thing. However, I'm hoping to extend the functionality (e.g. such that each row is uniquely transformed).

[jakevdp]

Good question – I think the issue is that your input to vjp_fun is not of the expected size. The VJP is a backward pass, so it expects an input with shape/size corresponding to the cotangent; since you've defined dJ_du is the same shape as y, this should work:

y, vjp_fun = vjp(mapping, x)
print(vjp_fun(dJ_du))

But if you vmap the vjp_fun before calling it on a one-dimensional input, it is the equivalent of passing a scalar to vjp_fun, which will fail.

Does that answer your question?

[smartalecH]

Aha, you're absolutely right, @jakevdp! Thanks for the quick reply.

Just for anybody else who stumbles upon this thread hoping to resolve an InconclusiveDimensionOperation error (I was surprised to not see any existing issues/discussions) you can still use a vmap, just ensure the input has the correct dimensionality (as @jakevdp points out).

In other words, this also works:

dJ_du = 0.5*npj.ones((1,N1*N0))
mg, = vmap(vjp_fun)(dJ_du)

As does this (which emulates a tensor jacbian product):

dJ_du = 0.5*npj.ones((10,N1*N0))
mg, = vmap(vjp_fun)(dJ_du)

Seems I forgot about that little nuance.

Thanks again for the help!