transpose over batch of 3d arrays using vmap

[mgbukov]

I have a batch of B 3D arrays data. Each 3D array can be transposed according to some fixed, but different axis defined in transposed_axes. I am trying to use vmap to push the loop over the batch out of python into cpp for speed.

Here's my attempt:

from jax import vmap
import numpy as np
from itertools import permutations 


seed=0
np.random.seed(0)

# define 3D arrays
B=5
data = np.random.uniform(size=(B,2,2,2) )

# define axes to transpose: each array has a different transpose axis. 
all_transposed_axes = np.array( list(permutations(range(3), )), dtype=np.int8 )
random_inds = np.random.choice(np.arange(all_transposed_axes.shape[0]), size=B)
transposed_axes = all_transposed_axes[random_inds]

# check shapes 
print(data.shape, transposed_axes.shape) # (5, 2, 2, 2) (5, 3)

# define function do transpose a single 3D array
def data_transpose(data,transposed_axes):
	return data.transpose(transposed_axes)

# check that function works on a single data points
result = data_transpose(data[0,:], transposed_axes[0,:]) # works

# use vmap to vectorize above function
vec_transpose = vmap(data_transpose,in_axes=(0,0), )

# call vectorized function
result = vec_transpose(data, transposed_axes) # throws error

This throws the error, as shown below. I think it occurs because jnp.transpose() only accepts lists and tuples for the argument axis, but not arrays, unlike np.transpose(). But I don't know how to define an array with tuples so that jax can correctly recognize the data types and does not complain.

After trying various things (like data.transpose( tuple(transposed_axes) ), and changing data types and array types, etc.), I am not sure if there is a way to make this work in jax -- any suggestions are most welcome!.

Traceback (most recent call last):
  File "test.py", line 90, in <module>
    result = vec_transpose(data, transposed_axes) # throws error
  File "test.py", line 82, in data_transpose
    return data.transpose(transposed_axes)
  File "/Users/mgbukov/miniconda3/envs/jax/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 1347, in _transpose
    axis = args[0] if args[0] is None else _ensure_index_tuple(args[0])
  File "/Users/mgbukov/miniconda3/envs/jax/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 1341, in _ensure_index_tuple
    return tuple(map(operator.index, x))
jax._src.traceback_util.FilteredStackTrace: jax._src.errors.TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int8[])>with<BatchTrace(level=1/0)>
  with val = DeviceArray([1, 0, 0, 0, 0], dtype=int8)
       batch_dim = 0 (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError)

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

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

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

Traceback (most recent call last):
  File "test.py", line 90, in <module>
    result = vec_transpose(data, transposed_axes) # throws error
  File "/Users/mgbukov/miniconda3/envs/jax/lib/python3.8/site-packages/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/Users/mgbukov/miniconda3/envs/jax/lib/python3.8/site-packages/jax/api.py", line 1223, in batched_fun
    out_flat = batching.batch(
  File "/Users/mgbukov/miniconda3/envs/jax/lib/python3.8/site-packages/jax/linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "test.py", line 82, in data_transpose
    return data.transpose(transposed_axes)
  File "/Users/mgbukov/miniconda3/envs/jax/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 1347, in _transpose
    axis = args[0] if args[0] is None else _ensure_index_tuple(args[0])
  File "/Users/mgbukov/miniconda3/envs/jax/lib/python3.8/site-packages/jax/_src/numpy/lax_numpy.py", line 1341, in _ensure_index_tuple
    return tuple(map(operator.index, x))
  File "/Users/mgbukov/miniconda3/envs/jax/lib/python3.8/site-packages/jax/core.py", line 466, in __index__
    raise TracerIntegerConversionError(self)
jax._src.errors.TracerIntegerConversionError: The __index__() method was called on the JAX Tracer object Traced<ShapedArray(int8[])>with<BatchTrace(level=1/0)>
  with val = DeviceArray([1, 0, 0, 0, 0], dtype=int8)
       batch_dim = 0 (https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerIntegerConversionError)

[jakevdp]

The axes argument to jnp.transpose is required to be static, so it is not possible to vmap over this argument. This requirement of static transpose axes comes from XLA, so it is not something that can be easily modifed. You'd have to come up with some alternate way to transpose the array that can make use of dynamic inputs.

Here' an approach that might be sufficient for your needs:

def dynamic_transpose(x, axes):
  axes = jnp.asarray(axes)
  assert len(set(x.shape)) == 1
  assert axes.shape == (x.ndim,)
  indices = jnp.mgrid[tuple(slice(s) for s in x.shape)]
  indices = indices[axes]
  return x[tuple(indices[i] for i in range(indices.shape[0]))]

result_slow = jnp.array([x.transpose(p) for x, p in zip(data, transposed_axes)])
result_vmap = vmap(dynamic_transpose)(data, transposed_axes)
np.testing.assert_allclose(result_slow, result_vmap)

[mgbukov]

Thanks! This does the job indeed. In fact, one can jit dynamic_transpose for extra speed.

A quick follow-up Q: do the assert statements, being extra operations, slow-down the code under vmap? I guess one can release them in the end for extra speed.

[jakevdp]

The assert statements will run only at trace-time on the abstract values of the arrays passed to vmap, so I don't anticipate that they'd play a significant role in the computation time.