Unstack operation in Jax

[conceptofmind]

Hello,

In PyTorch, we have torch.unbind().

For example:

x1, x2 = x.unbind(dim=-2)

And, in Tensorflow, we have tf.unstack().

For example:

x1, x2 = x.unstack(dim=-2)

Would the equivalent in JAX be?

def jax_unstack(x, axis=0):
    return jnp.moveaxis(x, axis, 0)

x1, x2 = jax_unstack(x, axis =-2)

Any help would be greatly appreciated.

Thank you,

Enrico

[jakevdp]

Yes, that seems like a concise and performant way to implement this functionality in JAX. You might avoid the reshape by using lax.index_in_dim instead:

import jax.numpy as jnp
from jax import lax

def jax_unstack(x, axis=0):
  return [lax.index_in_dim(x, i, axis, keepdims=False) for i in range(x.shape[axis])]

x = jnp.arange(9).reshape(3, 3)
jax_unstack(x, axis=0)
# [DeviceArray([0, 1, 2], dtype=int32),
#  DeviceArray([3, 4, 5], dtype=int32),
#  DeviceArray([6, 7, 8], dtype=int32)]

[conceptofmind]

Hi @jakevdp ,

I greatly appreciate the help.

I am working on a rotary positional embedding. I am using einops to rearrange the shape of x before the unstack operation.

def jax_unstack(x, axis = 0):
    return jnp.moveaxis(x, axis, 0)

def rotate_half(x):
    x = rearrange(x, "... (j d) -> ... j d", j = 2)
    x1, x2 = jax_unstack(x, dim = -2)
    return jnp.concatenate((-x2, x1), dim = -1)

def apply_rotary_pos_emb(pos, t):
    return (t * pos.cos()) + (rotate_half(t) * pos.sin())

I will look into using your implementation to avoid having to do a reshape.

Thank you,

Enrico