lax.dynamic_slice: TypeError: broadcast_in_dim broadcast_dimensions must have length equal to operand ndim; got broadcast_dimensions (0,) for operand ndim 2.

[bhchiang]

Running into an obscure error when trying to use lax.dynamic_slice: TypeError: broadcast_in_dim broadcast_dimensions must have length equal to operand ndim; got broadcast_dimensions (0,) for operand ndim 2.

import jax
from jax import lax, vmap
import jax.numpy as jnp

def ixs(y_start, y_end, x_start, x_end):
    return jnp.meshgrid(jnp.arange(x_start, x_end), jnp.arange(y_start, y_end))


@jax.partial(jax.jit, static_argnums=(1, 2))
def f(img, search_window_radius, filter_radius):
    h, w = img.shape
    pad = search_window_radius
    img_pad = jnp.pad(img, pad)

    def compute(y, x):
        # (y, x) are the center of the current neighborhood
        center_patch = jax.lax.dynamic_slice(img_pad, (y-filter_radius, x-filter_radius), (2*filter_radius+1, 2*filter_radius+1, ))
        return jnp.mean(center_patch)

    out = vmap(compute)(*ixs(pad, h+pad, pad, w+pad))
    return out
  
f(jnp.arange(25).reshape(5, 5), 5, 2)

filter_radius is static, y and x are not which is permitted according to the lax.dynamic_slice documentation. img_pad is 2D and I have tuples of length 2 for start_indices and start_sizes. I'm not sure what the broadcast_dimensions = (0,) is referring to?

Short reproduction notebook here.

Thanks!

[bhchiang]

Another quirk I noticed is that when I return 1 inside compute(...), the returned shape is (5,) vs. the (15, 15) that I'm mapping over with vmap, which I've never run into before.

[bhchiang]

My mistake - needed to do another vmap within the first one.

[di-xiao-94]

Hi, I'm running into the exact same error as you while using dynamic_slice: "TypeError: broadcast_in_dim broadcast_dimensions must have length equal to operand ndim; got broadcast_dimensions (0,) for operand ndim 2.". I was wondering if you could clarify what you meant by needing to another vmap within the first one

[wbrenton]

Would be stellar if someone could link a fix for the attached notebook. Running into the same problem.