help with this small code: Can I sample from a list inside vmap?

[externalsupplierstaff]

What is the correct way to do the following?

idxs = jnp.array([1, 0])
def sample_from_list(idx, mylist):
  return mylist[idx]

print(idxs)
print(sample_from_list(idxs[0], list('abcd')))
jax.vmap(sample_from_list, (0, None), 0)(idxs, list('abcd'))

[externalsupplierstaff]

my workaround:

def my_vmap_first_arg(fn, first_arg, *args):
  ress = [fn(first_arg_e, *args) for first_arg_e in first_arg]
  return tuple(zip(*ress))[0]

because the original code fails with the following output:

[1 0]
b

---------------------------------------------------------------------------

UnfilteredStackTrace                      Traceback (most recent call last)

<ipython-input-59-424276c1c88f> in <module>()
     18 print(sample_from_list(idxs[0], list('abcd')))
---> 19 jax.vmap(sample_from_list, (0, None), 0)(idxs, list('abcd'))
     20 

6 frames

/usr/local/lib/python3.7/dist-packages/jax/_src/traceback_util.py in reraise_with_filtered_traceback(*args, **kwargs)
    182     try:
--> 183       return fun(*args, **kwargs)
    184     except Exception as e:

/usr/local/lib/python3.7/dist-packages/jax/_src/api.py in batched_fun(*args, **kwargs)
   1316         lambda: flatten_axes("vmap out_axes", out_tree(), out_axes)
-> 1317     ).call_wrapped(*args_flat)
   1318     return tree_unflatten(out_tree(), out_flat)

/usr/local/lib/python3.7/dist-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    165     try:
--> 166       ans = self.f(*args, **dict(self.params, **kwargs))
    167     except:

<ipython-input-59-424276c1c88f> in sample_from_list(idx, mylist)
     14 def sample_from_list(idx, mylist):
---> 15   return mylist[idx]
     16 

/usr/local/lib/python3.7/dist-packages/jax/core.py in __index__(self)
    471   def __index__(self):
--> 472     raise TracerIntegerConversionError(self)
    473 

UnfilteredStackTrace: jax._src.errors.TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int32[])>with<BatchTrace(level=1/0)>
  with val = DeviceArray([1, 0], dtype=int32)
       batch_dim = 0
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------


The above exception was the direct cause of the following exception:

TracerIntegerConversionError              Traceback (most recent call last)

<ipython-input-59-424276c1c88f> in <module>()
     17 print(idxs)
     18 print(sample_from_list(idxs[0], list('abcd')))
---> 19 jax.vmap(sample_from_list, (0, None), 0)(idxs, list('abcd'))
     20 
     21 # obss, rews, dones, env_states = jax.vmap(environment.reset, (0, None), 0)(idxs.reshape(-1, 1), list_of_images)

<ipython-input-59-424276c1c88f> in sample_from_list(idx, mylist)
     13 idxs = jnp.array([1, 0])
     14 def sample_from_list(idx, mylist):
---> 15   return mylist[idx]
     16 
     17 print(idxs)

TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int32[])>with<BatchTrace(level=1/0)>
  with val = DeviceArray([1, 0], dtype=int32)
       batch_dim = 0
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError

[jakevdp]

Hi - thanks for the question! This is not possible because JAX operations can only operate on JAX arrays (not lists), and JAX does not support a string dtype. If you had numerical data it would be fine:

jax.vmap(sample_from_list, (0, None), 0)(idxs, jnp.array([1, 2, 3, 4]))
# DeviceArray([2, 1], dtype=int32)

but there is no way to do such computation with lists of strings in JAX.

[externalsupplierstaff]

Thanks for the answer!