Pallas slicing Vector vs Matrix outputs different behaviour

[bsaoptima]

Hi there,

I noticed something strange in my code, especially when I handle slicing for a vector vs a matrix.

def add_vectors_kernel(x_ref, y_ref, start, end, o_ref):
    x, y = x_ref[start:end], y_ref[start:end]
    o_ref[start:end] = x * y

@partial(jax.jit, static_argnums=(2,3))
def add_vectors(x: jax.Array, y: jax.Array, start: int, end: int) -> jax.Array:
    return pl.pallas_call(
        lambda x,y,o_ref : add_vectors_kernel(x, y, start, end, o_ref),
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype)
    )(x, y)
add_vectors(jnp.arange(8), jnp.arange(8), 2, 6)

This simple code takes 2 vectors, slices them and then does element-wise multiplication. The output as expected is another vector full of 0 except where we placed the outputs.

Array([ 0,  0,  4,  9, 16, 25,  0,  0], dtype=int32)

If I do the same thing with matrices:

def add_vectors_kernel(x_ref, y_ref, start, end, o_ref):
    x, y = x_ref[:, start:end], y_ref[:, start:end]
    o_ref[:, start:end] = x * y

@partial(jax.jit, static_argnums=(2,3))
def add_vectors(x: jax.Array, y: jax.Array, start: int, end: int) -> jax.Array:
    return pl.pallas_call(
        lambda x,y,o_ref : add_vectors_kernel(x, y, start, end, o_ref),
        out_shape=jax.ShapeDtypeStruct(x.shape, x.dtype)
    )(x, y)
add_vectors(jnp.array([[1,2,3,4], [1,2,3,4], [1,2,3,4], [1,2,3,4]]), jnp.array([[1,2,3,4], [1,2,3,4], [1,2,3,4], [1,2,3,4]]), 0, 2)

Array([[1, 4, 3, 4],
       [1, 4, 3, 4],
       [1, 4, 3, 4],
       [1, 4, 3, 4]], dtype=int32)

This is strange because I expected the last 2 columns to be filled with 0 not with original values.

Anyone has an idea why this is happening?