def testStackArrayArgument(self):

    @api.jit
    def foo(x):
        return jnp.stack(x)
    foo(onp.zeros(2))

    @api.jit
    def foo(x):
        return jnp.concatenate(x)
    foo(onp.zeros((2, 2)))
def testStackArrayArgument(self):

    @api.jit
    def foo(x):
        return jnp.stack(x)
    foo(onp.zeros(2))

    @api.jit
    def foo(x):
        return jnp.concatenate(x)
    foo(onp.zeros((2, 2)))