vmap: eval an array without a primitive

[LeonEricsson]

help me understand what's going wrong inside this vectorized function.

def fun(idx, input)
    # the following line gives "ValueError: [eval] Attempting to eval an array without a primitive."
    x = max(0, idx)

    # the following line gives "ValueError: Slice indices must be integers or None."
    y = input[idx:]

v_fun = mx.vmap(fun, in_axes=(0, None))

input = mx.ones((20,))
idxs = mx.arange(1, 20)

out = v_fun(idxs, input_ids, ngram)

I might be mistaken but in my limited time using JAX I have not come across this issue. As I understand it idx becomes an array
of size 1, but can I not access this value?

[awni]

You are doing a couple things strangely that are likely causing issues:

1. Using Python max instead of mx.maximum
2. idx is an mx.array and not supported with Python style slicing (even scalar arrays won't work there). I'm not sure if we will add that, it's possible we could but for now, best to avoid it.

[LeonEricsson]

hmm

1. interesting, I was actually using mx.max but changing to mx.maximum did indeed solve that issue
2. changing the line to y = input[5:] works great? but perhaps I can't convert idx to a scalar? i'm allowed to do y = input[idx] but y = input[idx:] fails, so I guess it's slicing with a scalar array that doesn't work

[awni]

Yea exactly, you can't do a Python slice object like slice(x, y, z) with mx.arrays of any dimension. So any syntax which uses slice (like the [5:]) won't work with arrays.