Repository: google/jax

Issue:
np.stack should support axis argument



Discussion:

Resolving patch:
diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -995,6 +995,18 @@ def cumulative_reduction(a, axis=None, dtype=None):
 
 ### Array-creation functions
 
+# TODO(phawkins): use this helper everywhere.
+def _canonicalize_axis(axis, num_dims):
+  """Canonicalize an axis in (-num_dims, num_dims) to [0, num_dims)."""
+  if axis < 0:
+    axis = axis + num_dims
+  if axis < 0 or axis >= num_dims:
+      raise ValueError(
+          "axis {} is out of bounds for array of dimension {}".format(
+              axis, num_dims))
+  return axis
+
+
 @_wraps(onp.pad)
 def pad(array, pad_width, mode, constant_values=0):
   if mode != "constant":
@@ -1014,11 +1026,19 @@ def pad(array, pad_width, mode, constant_values=0):
 
 
 @_wraps(onp.stack)
-def stack(arrays):
+def stack(arrays, axis=0):
   if not arrays:
     raise ValueError("Need at least one array to stack.")
-  new_arrays = [reshape(x, (-1,) + onp.shape(x)) for x in arrays]
-  return reshape(concatenate(new_arrays), (len(arrays),) + arrays[0].shape)
+  shape0 = shape(arrays[0])
+  axis = _canonicalize_axis(axis, len(shape0) + 1)
+  new_shape = list(shape0)
+  new_shape.insert(axis, 1)
+  new_arrays = []
+  for a in arrays:
+    if shape(a) != shape0:
+      raise ValueError("All input arrays must have the same shape.")
+    new_arrays.append(reshape(a, new_shape))
+  return concatenate(new_arrays, axis=axis)
 
 
 @_wraps(onp.concatenate)
@@ -1044,6 +1064,11 @@ def hstack(tup):
   return concatenate(arrs, 1)
 
 
+@_wraps(onp.dstack)
+def dstack(tup):
+  return concatenate([atleast_3d(m) for m in tup], axis=2)
+
+
 @_wraps(onp.column_stack)
 def column_stack(tup):
   arrays = []
@@ -1073,6 +1098,19 @@ def atleast_2d(*arys):
     return [atleast_2d(arr) for arr in arys]
 
 
+@_wraps(onp.atleast_3d)
+def atleast_3d(*arys):
+  if len(arys) == 1:
+    arr = array(arys[0])
+    if ndim(arr) <= 1:
+      arr = arr.reshape((1, -1, 1))
+    elif ndim(arr) == 2:
+      arr = arr.reshape(shape(arr) + (1,))
+    return arr
+  else:
+    return [atleast_3d(arr) for arr in arys]
+
+
 # TODO(mattjj): can this be simplified?
 @_wraps(onp.array)
 def array(object, dtype=None, copy=True, order="K", ndmin=0):
@@ -1711,18 +1749,6 @@ def argsort(a, axis=-1, kind='quicksort', order=None):
     return perm
 
 
-# TODO(phawkins): use this helper everywhere.
-def _canonicalize_axis(axis, num_dims):
-  """Canonicalize an axis in (-num_dims, num_dims) to [0, num_dims)."""
-  if axis < 0:
-    axis = axis + num_dims
-  if axis < 0 or axis >= num_dims:
-      raise ValueError(
-          "axis {} is out of bounds for array of dimension {}".format(
-              axis, num_dims))
-  return axis
-
-
 @_wraps(onp.take)
 def take(a, indices, axis=None, out=None, mode=None):
   if 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
@@ -18,6 +18,7 @@
 
 import collections
 import functools
+from functools import partial
 import itertools
 from unittest import skip
 
@@ -108,6 +109,9 @@ def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None):
 
 JAX_COMPOUND_OP_RECORDS = [
     op_record("angle", 1, number_dtypes, all_shapes, jtu.rand_default(), []),
+    op_record("atleast_1d", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
+    op_record("atleast_2d", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
+    op_record("atleast_3d", 1, default_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("conjugate", 1, number_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
     op_record("deg2rad", 1, float_dtypes, all_shapes, jtu.rand_default(), []),
     op_record("divide", 2, number_dtypes, all_shapes, jtu.rand_nonzero(), ["rev"]),
@@ -664,9 +668,9 @@ def testTrace(self, shape, dtype, out_dtype, offset, axis1, axis2, rng):
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
 
   @parameterized.named_parameters(jtu.cases_from_list(
-      {"testcase_name": "_{}".format(
-          jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes)),
-       "shape": shape, "dtypes": dtypes, "rng": rng}
+      {"testcase_name": "_{}_axis={}".format(
+          jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes), axis),
+       "shape": shape, "axis": axis, "dtypes": dtypes, "rng": rng}
       for dtypes in [
         [onp.float32],
         [onp.float32, onp.float32],
@@ -675,10 +679,35 @@ def testTrace(self, shape, dtype, out_dtype, offset, axis1, axis2, rng):
         [onp.float32, onp.int32, onp.float64],
       ]
       for shape in [(), (2,), (3, 4), (1, 100)]
+      for axis in range(-len(shape), len(shape) + 1)
       for rng in [jtu.rand_default()]))
-  def testStack(self, shape, dtypes, rng):
+  def testStack(self, shape, axis, dtypes, rng):
     args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
-    self._CheckAgainstNumpy(lnp.stack, onp.stack, args_maker, check_dtypes=True)
+    onp_fun = partial(onp.stack, axis=axis)
+    lnp_fun = partial(lnp.stack, axis=axis)
+    self._CheckAgainstNumpy(lnp_fun, onp_fun, args_maker, check_dtypes=True)
+
+
+  @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, "rng": rng}
+      for op in ["hstack", "vstack", "dstack"]
+      for dtypes in [
+        [onp.float32],
+        [onp.float32, onp.float32],
+        [onp.float32, onp.int32, onp.float32],
+        [onp.float32, onp.int64, onp.float32],
+        [onp.float32, onp.int32, onp.float64],
+      ]
+      for shape in [(), (2,), (3, 4), (1, 100), (2, 3, 4)]
+      for rng in [jtu.rand_default()]))
+  def testHVDStack(self, shape, op, dtypes, rng):
+    args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
+    onp_fun = getattr(onp, op)
+    lnp_fun = getattr(lnp, op)
+    self._CheckAgainstNumpy(lnp_fun, onp_fun, args_maker, check_dtypes=True)
+
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_inshape={}_outdtype={}".format(