How could I change this function to be applicable by jax jit?

[alisheikholeslam]

let say that I have:

a = np.array([[ 1,2,3],
              [ 4,5,6],
              [ 7,8,9]])

I want to apply jax jit on the following function:

def ex(a, N):
    out = np.zeros( (N+1)*np.array(a.shape)-N,dtype=a.dtype)
    out[::N+1,::N+1] = a
    return out

I tried:

@jax.jit
def new(a, N):
    out = jnp.zeros( (N+1)*jnp.array(a.shape)-N,dtype=a.dtype)
    out = out.at[::N+1,::N+1].set(a)
    return out

which will stuck by an error:

TypeError: Shapes must be 1D sequences of concrete values of integer type, got (Traced<ShapedArray(int64[2])>with<DynamicJaxprTrace(level=0/1)>,).
If using jit, try using static_argnums or applying jit to smaller subfunctions.

How must I use static_argnums=1? (@partial(jax.jit, static_argnums=1) will not work)
I am stuck in creating out array and can not check the second line whether it works or not. Does it need to be changed?

[YouJiacheng]

I think following code will work.

@partial(jax.jit, static_argnums=1)
def new(a, N):
    out = jnp.zeros( (N+1)*np.array(a.shape)-N,dtype=a.dtype) # note: use np for shape, not jnp
    out = out.at[::N+1,::N+1].set(a)
    return out

[alisheikholeslam]

is there any equivalent to use instead of (N+1)*np.array(a.shape)-N if a be a jnp.array that make better the performance? i.e. if

b = jnp.array(a)
(N+1)*np.array(a.shape)-N == (N+1)*jnp.array(jnp.shape(b))-N

[YouJiacheng]

np.array(a.shape) generally has a better performance than jnp.array(a.shape), since the a.shape as an array is usally very small.
And IIRC, the best practice of JAX suggest that you should use np for static value(i.e. compile-time constant), e.g. shape of array. Most of internal functions use np for shape operations.