Repository: google/jax

Issue:
hstack and vstack produce very inefficient jaxpr and jit slowly; possible fix with reshape?
hstack is very inefficient for tensors as it produces jaxpr code with length proportional to size of the traced array.

Compare:
```
{ lambda  ; a b c.
  let d = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4)
                            shape=(1, 2, 2, 3, 3) ] a
      e = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4)
                            shape=(1, 2, 2, 3, 3) ] b
      f = concatenate[ dimension=0 ] d e
      g = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4)
                            shape=(1, 2, 2, 3, 3) ] b
      h = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4)
                            shape=(1, 2, 2, 3, 3) ] c
      i = concatenate[ dimension=0 ] g h
      j = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4, 5)
                            shape=(1, 2, 2, 2, 3, 3) ] f
      k = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4, 5)
                            shape=(1, 2, 2, 2, 3, 3) ] i
      l = concatenate[ dimension=0 ] j k
      m = slice[ limit_indices=(1, 2, 2, 2, 3, 3)
                 start_indices=(0, 0, 0, 0, 0, 0)
                 strides=(1, 1, 1, 1, 1, 1) ] l
      n = squeeze[ dimensions=(0,) ] m
      o = slice[ limit_indices=(2, 2, 2, 2, 3, 3)
                 start_indices=(1, 0, 0, 0, 0, 0)
                 strides=(1, 1, 1, 1, 1, 1) ] l
      p = squeeze[ dimensions=(0,) ] o
      q = concatenate[ dimension=1 ] n p
      r = slice[ limit_indices=(1, 4, 2, 3, 3)
                 start_indices=(0, 0, 0, 0, 0)
                 strides=(1, 1, 1, 1, 1) ] q
      s = squeeze[ dimensions=(0,) ] r
      t = slice[ limit_indices=(2, 4, 2, 3, 3)
                 start_indices=(1, 0, 0, 0, 0)
                 strides=(1, 1, 1, 1, 1) ] q
      u = squeeze[ dimensions=(0,) ] t
      v = concatenate[ dimension=1 ] s u
      w = slice[ limit_indices=(1, 4, 3, 3)
                 start_indices=(0, 0, 0, 0)
                 strides=(1, 1, 1, 1) ] v
      x = squeeze[ dimensions=(0,) ] w
      y = slice[ limit_indices=(2, 4, 3, 3)
                 start_indices=(1, 0, 0, 0)
                 strides=(1, 1, 1, 1) ] v
      z = squeeze[ dimensions=(0,) ] y
      ba = slice[ limit_indices=(3, 4, 3, 3)
                  start_indices=(2, 0, 0, 0)
                  strides=(1, 1, 1, 1) ] v
      bb = squeeze[ dimensions=(0,) ] ba
      bc = slice[ limit_indices=(4, 4, 3, 3)
                  start_indices=(3, 0, 0, 0)
                  strides=(1, 1, 1, 1) ] v
      bd = squeeze[ dimensions=(0,) ] bc
      be = concatenate[ dimension=1 ] x z bb bd
      bf = slice[ limit_indices=(1, 12, 3)
                  start_indices=(0, 0, 0)
                  strides=(1, 1, 1) ] be
      bg = squeeze[ dimensions=(0,) ] bf
      bh = slice[ limit_indices=(2, 12, 3)
                  start_indices=(1, 0, 0)
                  strides=(1, 1, 1) ] be
      bi = squeeze[ dimensions=(0,) ] bh
      bj = slice[ limit_indices=(3, 12, 3)
                  start_indices=(2, 0, 0)
                  strides=(1, 1, 1) ] be
      bk = squeeze[ dimensions=(0,) ] bj
      bl = slice[ limit_indices=(4, 12, 3)
                  start_indices=(3, 0, 0)
                  strides=(1, 1, 1) ] be
      bm = squeeze[ dimensions=(0,) ] bl
      bn = concatenate[ dimension=1 ] bg bi bk bm
  in (bn,) }
```

to a better, equivalent code that can be achieved using jnp.reshape

```
{ lambda  ; a b c.
  let d = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4)
                            shape=(1, 2, 2, 3, 3) ] a
      e = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4)
                            shape=(1, 2, 2, 3, 3) ] b
      f = concatenate[ dimension=0 ] d e
      g = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4)
                            shape=(1, 2, 2, 3, 3) ] b
      h = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4)
                            shape=(1, 2, 2, 3, 3) ] c
      i = concatenate[ dimension=0 ] g h
      j = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4, 5)
                            shape=(1, 2, 2, 2, 3, 3) ] f
      k = broadcast_in_dim[ broadcast_dimensions=(1, 2, 3, 4, 5)
                            shape=(1, 2, 2, 2, 3, 3) ] i
      l = concatenate[ dimension=0 ] j k
      m = reshape[ dimensions=(0, 2, 4, 1, 3, 5)
                   new_sizes=(12, 12) ] l
  in (m,) }
```

----

Probably `hstack` can be re-expressed in terms of reshape in general. I'm new to `jax` so maybe there are some negative side effects to such approach?

----

Code to reproduce issue:

```python
import jax
import jax.numpy as jnp

n = 2

mAA = 1.0*jnp.arange(3*n*3*n).reshape((n,n,3,3))
mBB = 10.0*jnp.arange(3*n*3*n).reshape((n,n,3,3))
mAB = 2.0*jnp.arange(3*n*3*n).reshape((n,n,3,3))

def stack_hard(AA,AB,BB):
    return jnp.hstack(
        jnp.hstack(
            jnp.hstack(
                jnp.hstack(
                    jnp.array(
                        [[AA,AB],[AB,BB]]
                    )
                )
            )
        )
    )

def stack_easy(AA,AB,BB):
    return  jax.lax.reshape(
                jnp.array([[AA,AB],[AB,BB]]),
                (6*n,6*n),
                dimensions = (0,2,4,1,3,5)
            )

# JIT is very slow in case of larger n
# fast_stack = jax.jit(stack_hard)
# fast_stack(mAA,mBB,mAB)

print('===========================')
print(
    jax.make_jaxpr(stack_hard)(mAA,mAB,mBB)
    )

print('===========================')
print(
    jax.make_jaxpr(stack_easy)(mAA,mAB,mBB)
    )


print(stack_easy(mAA,mAB,mBB))
print(stack_hard(mAA,mAB,mBB))
```



Discussion:
Hi - the issue here is that the call signature of `hstack` is that it accepts a single argument, which is a tuple of arrays.

A tuple is a Python concept, not an XLA concept, so when you pass an array to something that expects a tuple, it must be converted into `N` array objects that are then passed back to XLA.

I'm not sure what we could do to "fix" this – maybe we could raise an error in the case that a single array is passed to `hstack`, to prevent this sort of silent conversion back to a numpy tuple, and require users to pass `tuple(arr)` explicitly. It would be less convenient, but it would make more apparent the computational cost implicit in the function's signature.

What do you think?
Also, I don't think it's generally true that `hstack` of a single array can be expressed in terms of a reshape. Here's a counter-example:
```python
>>> import jax.numpy as jnp
>>> x = jnp.arange(12).reshape(3, 2, 2)
>>> jnp.hstack(x)
DeviceArray([[ 0,  1,  4,  5,  8,  9],
             [ 2,  3,  6,  7, 10, 11]], dtype=int32)
```
I'm not sure if there's any alternative here other than to split the array into three and pass them to ``lax.concat``, which is what `hstack` currently does.
Oh, I understand better why this happened. Perhaps we can improve just the case where only one `jnp` array is passed as argument.

I'm pretty sure in all cases `jnp.hstack` can be expressed with `jax.lax.reshape` (note: *not* `jnp.reshape`) due to it's cool feature of optional arg `dimensions`.

In case of your example it would be:

``` python
>>> import jax numpy as jnp
>>> import jax
>>> x = jnp.arange(12).reshape(3, 2, 2)
>>> jax.lax.reshape(x,(2,6),dimensions=(1,0,2)) - jnp.hstack(x)
DeviceArray([[0, 0, 0, 0, 0, 0],
             [0, 0, 0, 0, 0, 0]], dtype=int32)
```
Nice, I didn't know about that! If it's not too complicated, I think special-casing the single array case in terms of these kinds of reshapes would be worthwhile - maybe it could all be handled at the level of `jnp.concatenate`
> Hi - the issue here is that the call signature of `hstack` is that it accepts a single argument, which is a tuple of arrays.
> 
> A tuple is a Python concept, not an XLA concept, so when you pass an array to something that expects a tuple, it must be converted into `N` array objects that are then passed back to XLA.
> 
> I'm not sure what we could do to "fix" this – maybe we could raise an error in the case that a single array is passed to `hstack`, to prevent this sort of silent conversion back to a numpy tuple, and require users to pass `tuple(arr)` explicitly. It would be less convenient, but it would make more apparent the computational cost implicit in the function's signature.
> 
> What do you think?

Is there a way to see which operations in a program might be undergoing such a forced conversion from XLA to numpy back to XLA? I am currently suffering a significant slow down in a program I have written and I am wondering if it due to similar issues going on in the background.
The main source of this kind of thing is calling `__iter__` on a `DeviceArray`.
> The main source of this kind of thing is calling `__iter__` on a `DeviceArray`.

How can I find out if this is the case?
Alternatively, I think it might make sense to just raise an error when a single array is passed into `hstack`/`vstack`/`stack`/`concatenate`. We would force the user to write something more explicit like `jnp.vstack(list(array))`, which hopefully has more obvious performance implications.
That’s two votes for erroring on implicit tuple conversion - maybe that’s a cleaner route
> That’s two votes for erroring on implicit tuple conversion - maybe that’s a cleaner route

Oh, I see you did suggest that already 👍 
Erroring implicit tuple conversion :+1:

As for `hstack` and `vstack` I'd still prefer if there was a special case with reshape, especially since the workaround is not immediately obvious.

I think I have working code for this (up to edge cases such as passing a scalar etc.)

``` python
import jax.numpy as jnp
import jax
import numpy as np


def hstack_alternative(array):
    shp = array.shape
    return jax.lax.reshape(
        array,
        (shp[1], shp[0] * shp[2]) + shp[3:],
        dimensions=((1, 0) + tuple(range(2, len(shp)))),
    )


def vstack_alternative(array):
    shp = array.shape
    return jax.lax.reshape(array, (shp[0] * shp[1],) + shp[2:])


x = jnp.arange(12).reshape(3, 2, 2)
y = jnp.arange(2 * 3 * 5).reshape(2, 3, 5)
z = jnp.arange(2 * 3 * 5 * 7).reshape(2, 3, 5, 7)


def test_hstack_alternative():
    assert np.allclose(jnp.hstack(x), hstack_alternative(x))
    assert np.allclose(jnp.hstack(y), hstack_alternative(y))
    assert np.allclose(jnp.hstack(z), hstack_alternative(z))


def test_vstack_alternative():
    assert np.allclose(jnp.vstack(x), vstack_alternative(x))
    assert np.allclose(jnp.vstack(y), vstack_alternative(y))
    assert np.allclose(jnp.vstack(z), vstack_alternative(z))
```

That's very cool – it would be worth adding this code-path to `vstack` and `hstack` in my opinion.

Are you interested in putting together a PR?
I was playing with this a bit - here's the implementation of `jnp.concatenate` for array inputs in terms of `lax.reshape`:
```python
import jax.numpy as jnp
from jax import lax
from jax._src.util import canonicalize_axis

def _concatenate(x, axis=0):
  assert isinstance(x, jnp.ndarray)
  if x.ndim == 0:
    raise ValueError("Need at least one array to concatenate.")
  if x.ndim == 1:
    raise ValueError("Zero-dimensional arrays cannot be concatenated.")
  axis = canonicalize_axis(axis, x.ndim - 1)
  shape = x.shape[1:axis + 1] + (x.shape[0] * x.shape[axis + 1],) + x.shape[axis + 2:]
  dimensions = [*range(1, axis + 1), 0, *range(axis + 1, x.ndim)]
  return lax.reshape(x, shape, dimensions)
```
Quickly tested with:
```python
import numpy as np

x = jnp.arange(2*3*4*5*4*3*2).reshape(2, 3, 4, 5, 4, 3, 2)
for axis in range(1 - x.ndim, x.ndim - 1):
  c1 = jnp.concatenate(x, axis=axis)
  c2 = _concatenate(x, axis=axis)
  np.testing.assert_array_equal(c1, c2)
```
The implementation is simple enough that I think we should add it to JAX, along with similar approaches for `hstack`, `vstack`, and perhaps other related functions.
> That's very cool – it would be worth adding this code-path to `vstack` and `hstack` in my opinion.
> 
> Are you interested in putting together a PR?

Sure :)
I'm new here so if you could help with the github side of things I'll be thankful.
Great! We have a bit of contribution information here: https://jax.readthedocs.io/en/latest/contributing.html#contributing-code-using-pull-requests

Feel free to open a work-in-progress PR if that would be helpful, and let me know if you have any questions
Hi @RadostW - just checking in. Is this still something you'd like to work on? If not, I can plan to put together the fix.

Either way, please let me know - thanks!

Resolving patch:
diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -2801,15 +2801,19 @@ def stack(arrays, axis: int =0, out=None):
     raise ValueError("Need at least one array to stack.")
   if out is not None:
     raise NotImplementedError("The 'out' argument to jnp.stack is not supported.")
-  _check_arraylike("stack", *arrays)
-  shape0 = shape(arrays[0])
-  axis = _canonicalize_axis(axis, len(shape0) + 1)
-  new_arrays = []
-  for a in arrays:
-    if shape(a) != shape0:
-      raise ValueError("All input arrays must have the same shape.")
-    new_arrays.append(expand_dims(a, axis))
-  return concatenate(new_arrays, axis=axis)
+  if isinstance(arrays, ndarray):
+    axis = _canonicalize_axis(axis, arrays.ndim)
+    return concatenate(expand_dims(arrays, axis + 1), axis=axis)
+  else:
+    _check_arraylike("stack", *arrays)
+    shape0 = shape(arrays[0])
+    axis = _canonicalize_axis(axis, len(shape0) + 1)
+    new_arrays = []
+    for a in arrays:
+      if shape(a) != shape0:
+        raise ValueError("All input arrays must have the same shape.")
+      new_arrays.append(expand_dims(a, axis))
+    return concatenate(new_arrays, axis=axis)
 
 @_wraps(np.tile)
 def tile(A, reps):
@@ -2868,32 +2872,41 @@ def concatenate(arrays, axis: int = 0):
 
 @_wraps(np.vstack)
 def vstack(tup):
-  return concatenate([atleast_2d(m) for m in tup], axis=0)
+  if isinstance(tup, ndarray):
+    arrs = jax.vmap(atleast_2d)(tup)
+  else:
+    arrs = [atleast_2d(m) for m in tup]
+  return concatenate(arrs, axis=0)
 row_stack = vstack
 
 
 @_wraps(np.hstack)
 def hstack(tup):
-  arrs = [atleast_1d(m) for m in tup]
-  if arrs[0].ndim == 1:
-    return concatenate(arrs, 0)
-  return concatenate(arrs, 1)
+  if isinstance(tup, ndarray):
+    arrs = jax.vmap(atleast_1d)(tup)
+    arr0_ndim = arrs.ndim - 1
+  else:
+    arrs = [atleast_1d(m) for m in tup]
+    arr0_ndim = arrs[0].ndim
+  return concatenate(arrs, axis=0 if arr0_ndim == 1 else 1)
 
 
 @_wraps(np.dstack)
 def dstack(tup):
-  return concatenate([atleast_3d(m) for m in tup], axis=2)
+  if isinstance(tup, ndarray):
+    arrs = jax.vmap(atleast_3d)(tup)
+  else:
+    arrs = [atleast_3d(m) for m in tup]
+  return concatenate(arrs, axis=2)
 
 
 @_wraps(np.column_stack)
 def column_stack(tup):
-  arrays = []
-  for v in tup:
-    arr = asarray(v)
-    if arr.ndim < 2:
-      arr = atleast_2d(arr).T
-    arrays.append(arr)
-  return concatenate(arrays, 1)
+  if isinstance(tup, ndarray):
+    arrs = jax.vmap(lambda x: atleast_2d(x).T)(tup) if tup.ndim < 3 else tup
+  else:
+    arrs = [atleast_2d(arr).T if arr.ndim < 2 else arr for arr in map(asarray, tup)]
+  return concatenate(arrs, 1)
 
 
 @_wraps(np.choose, skip_params=['out'])


Test patch:
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -2762,9 +2762,9 @@ def testDigitize(self, xshape, binshape, right, reverse, dtype):
     self._CompileAndCheck(jnp_fun, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_{}".format(
-          jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes)),
-       "shape": shape, "dtypes": dtypes}
+      {"testcase_name": "_{}_array={}".format(
+          jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes), array_input),
+       "shape": shape, "dtypes": dtypes, "array_input": array_input}
       for dtypes in [
         [np.float32],
         [np.float32, np.float32],
@@ -2772,19 +2772,23 @@ def testDigitize(self, xshape, binshape, right, reverse, dtype):
         [np.float32, np.int64, np.float32],
         [np.float32, np.int32, np.float64],
       ]
-      for shape in [(), (2,), (3, 4), (1, 5)]))
-  def testColumnStack(self, shape, dtypes):
+      for shape in [(), (2,), (3, 4), (1, 5)]
+      for array_input in [True, False]))
+  def testColumnStack(self, shape, dtypes, array_input):
     rng = jtu.rand_default(self.rng())
-    args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
+    if array_input:
+      args_maker = lambda: [np.array([rng(shape, dtype) for dtype in dtypes])]
+    else:
+      args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
     np_fun = _promote_like_jnp(np.column_stack)
     jnp_fun = jnp.column_stack
     self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_{}_axis={}".format(
-          jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes), axis),
-       "shape": shape, "axis": axis, "dtypes": dtypes}
+      {"testcase_name": "_{}_axis={}_array={}".format(
+          jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes), axis, array_input),
+       "shape": shape, "axis": axis, "dtypes": dtypes, "array_input": array_input}
       for dtypes in [
         [np.float32],
         [np.float32, np.float32],
@@ -2793,19 +2797,23 @@ def testColumnStack(self, shape, dtypes):
         [np.float32, np.int32, np.float64],
       ]
       for shape in [(), (2,), (3, 4), (1, 100)]
-      for axis in range(-len(shape), len(shape) + 1)))
-  def testStack(self, shape, axis, dtypes):
+      for axis in range(-len(shape), len(shape) + 1)
+      for array_input in [True, False]))
+  def testStack(self, shape, axis, dtypes, array_input):
     rng = jtu.rand_default(self.rng())
-    args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
+    if array_input:
+      args_maker = lambda: [np.array([rng(shape, dtype) for dtype in dtypes])]
+    else:
+      args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
     np_fun = _promote_like_jnp(partial(np.stack, axis=axis))
     jnp_fun = partial(jnp.stack, axis=axis)
     self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_op={}_{}".format(
-          op, jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes)),
-       "shape": shape, "op": op, "dtypes": dtypes}
+      {"testcase_name": "_op={}_{}_array={}".format(
+          op, jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes), array_input),
+       "shape": shape, "op": op, "dtypes": dtypes, "array_input": array_input}
       for op in ["hstack", "vstack", "dstack"]
       for dtypes in [
         [np.float32],
@@ -2814,10 +2822,14 @@ def testStack(self, shape, axis, dtypes):
         [np.float32, np.int64, np.float32],
         [np.float32, np.int32, np.float64],
       ]
-      for shape in [(), (2,), (3, 4), (1, 100), (2, 3, 4)]))
-  def testHVDStack(self, shape, op, dtypes):
+      for shape in [(), (2,), (3, 4), (1, 100), (2, 3, 4)]
+      for array_input in [True, False]))
+  def testHVDStack(self, shape, op, dtypes, array_input):
     rng = jtu.rand_default(self.rng())
-    args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
+    if array_input:
+      args_maker = lambda: [np.array([rng(shape, dtype) for dtype in dtypes])]
+    else:
+      args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
     np_fun = _promote_like_jnp(getattr(np, op))
     jnp_fun = getattr(jnp, op)
     self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)