Error on jax.numpy.tile when reps is a jax.numpy.array

[ethanleeman]

Not sure if this is a bug or WAI.

Consider the following 2 code examples that uses jnp.tile:

import jax
import jax.numpy as jnp
import numpy as np

@jax.jit
def tile_0():
  x = jnp.array([0, 1, 2])
  y = (2, 2)
  return jnp.tile(x, y)
tile_0()

@jax.jit
def tile_1():
  x = jnp.array([0, 1, 2])
  y = np.array([2, 2])
  return jnp.tile(x, y)
tile_1()

Both of these work with no errors. Now consider the following modification:

@jax.jit
def tile_2():
  x = jnp.array([0, 1, 2])
  y = jnp.array([2, 2])
  return jnp.tile(x, y)
tile_2()

This fails with the following error message:

AttributeError                            Traceback (most recent call last)
[google3/third_party/py/numpy/core/fromnumeric.py](https://colab.corp.google.com/drive/17m5KuUGKTtIM7JNKv8mXY-UWENaNryyz#) in ndim(a)
   3208     try:
-> 3209         return a.ndim
   3210     except AttributeError:

AttributeError: 'list' object has no attribute 'ndim'

During handling of the above exception, another exception occurred:

TracerArrayConversionError                Traceback (most recent call last)
5 frames
[google3/third_party/py/numpy/core/fromnumeric.py](https://colab.corp.google.com/drive/17m5KuUGKTtIM7JNKv8mXY-UWENaNryyz#) in ndim(a)
   3209         return a.ndim
   3210     except AttributeError:
-> 3211         return asarray(a).ndim
   3212 
   3213 

TracerArrayConversionError: The numpy.ndarray conversion method __array__() was called on traced array with shape int32[].
The error occurred while tracing the function tile_2 at <ipython-input-23-53638b53fd53>:1 for jit. This value became a tracer due to JAX operations on these lines:

  operation a:i32[2] = convert_element_type[new_dtype=int32 weak_type=False] b
    from line <ipython-input-23-53638b53fd53>:4:6 (tile_2)
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerArrayConversionError

Should jnp.tile fail in this instance? I am looking at the readthedoc https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.tile.html
and it seems like reps is just array_like so I am perplexed why this should fail.

Thank you!

[jakevdp]

This is expected: the reps argument must be a sequence of static integers (because the size of the output array is directly related to the values in reps). Within a JIT context, all JAX arrays are non-static, so you wouldn't be able to pass a JAX array to reps.

We need to update the docs here; reps is not array_like, but rather a sequence of integers. We'll get to that as part of #21461, but haven't done jax.numpy.tile yet.