How to reshape a single axis in an array

[magnusdk]

Hi!

I'm trying to write a function that takes an array and reshapes a single axis of that array into a new shape. E.g., we want to reshape an array of shape (5, 6, 7) into (5, 2, 3, 7), thus splitting the middle axis into two parts. It should work on arrays of any shape, given an axis and the new shape of the axis (assuming the axis can be reshaped to the new shape).

I want to do this in a jitted function, but am encountering a TracerIntegerConversionError. Here is code to reproduce:

import jax
import jax.numpy as jnp

@jax.jit
def reshape_axis(a: jnp.ndarray, axis: int, new_shape: tuple):
    "Reshape a single axis into a new shape."
    new_shape = a.shape[:axis] + new_shape + a.shape[axis + 1 :]
    return a.reshape(new_shape)

a = jnp.zeros((5, 6, 7))
new_a = reshape_axis(a, 1, (2, 3))
assert new_a.shape == (5, 2, 3, 7)

The problem is that a.shape[:axis] and a.shape[axis + 1 :] have dynamic sizes, dependent on the value of axis. However, once they are combined back into new_shape the size is deterministic.

I've tried to make the axis and new_shape arguments static, but when I then call it from another jitted function the program crashes because axis has type jax.interpreters.partial_eval.DynamicJaxprTracer. Reproduce:

import jax
import jax.numpy as jnp
from functools import partial

@partial(jax.jit, static_argnums=(1, 2), inline=True)
def reshape_axis(a: jnp.ndarray, axis: int, new_shape: tuple):
    "Reshape a single axis into a new shape."
    new_shape = a.shape[:axis] + new_shape + a.shape[axis + 1 :]
    return a.reshape(new_shape)

@jax.jit
def do_thing(a, axis, new_shape):
    return reshape_axis(a, axis, new_shape)

a = jnp.zeros((5, 6, 7))
assert do_thing(a, 1, (2, 3)).shape == (5, 2, 3, 7)

I could of course make the arguments static in do_thing as well, but it would be helpful if this implementation detail didn't leak out. I suspect there is a better way of implementing reshape_axis but I haven't found it yet. Any advice for how to proceed?

BTW, I've been using JAX at work for over 6 months now, and it's truly a pleasure. Thank you!

[jakevdp]

Hi - I think you've already landed on the correct solution: you need the static arguments to be marked static at the call-site. You could do this by having the same static_argnums annotation for do_thing that you have for reshape_axis.

[magnusdk]

Hi, thanks for the response! I also see that jax.numpy.swapaxes solves the problem in the same way.

In the previous iteration of the function we always reshape the last axis of an array, but I wanted something more general. I guess in that case the axis -1 is a static value, I just never thought about it that way until now.