Repository: google/jax

Issue:
add np.append and np.polyval
Requested by @duvenaud!


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
@@ -1070,6 +1070,27 @@ def diag(v, k=0):
     raise ValueError("diag input must be 1d or 2d")
 
 
+@_wraps(onp.polyval)
+def polyval(p, x):
+  if isinstance(p, onp.poly1d):
+    p = onp.asarray(p)
+  if isinstance(x, onp.poly1d):
+    y = 0
+  else:
+    y = zeros_like(x)
+  for i in range(len(p)):
+    y = y * x + p[i]
+  return y
+
+
+@_wraps(onp.append)
+def append(arr, values, axis=None):
+  if axis is None:
+    return concatenate([ravel(arr), ravel(values)], 0)
+  else:
+    return concatenate([arr, values], axis=axis)
+
+
 ### Tensor contraction operations
 
 


Test patch:
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -458,7 +458,7 @@ def wrapped_fun(*args):
   def _CheckAgainstNumpy(self, lax_op, numpy_reference_op, args_maker,
                          check_dtypes=False, tol=1e-5):
     args = args_maker()
-    lax_ans = lax_op(*args)
     numpy_ans = numpy_reference_op(*args)
+    lax_ans = lax_op(*args)
     self.assertAllClose(lax_ans, numpy_ans, check_dtypes=check_dtypes,
                         atol=tol, rtol=tol)
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
@@ -35,7 +35,8 @@
 config.parse_flags_with_absl()
 FLAGS = config.FLAGS
 
-nonempty_array_shapes = [(), (4,), (3, 4), (3, 1), (1, 4), (2, 1, 4), (2, 3, 4)]
+nonempty_nonscalar_array_shapes = [(4,), (3, 4), (3, 1), (1, 4), (2, 1, 4), (2, 3, 4)]
+nonempty_array_shapes = [()] + nonempty_nonscalar_array_shapes
 empty_array_shapes = [(0,), (0, 4), (3, 0),]
 
 scalar_shapes = [jtu.NUMPY_SCALAR_SHAPE]
@@ -119,6 +120,7 @@ def op_record(name, nargs, dtypes, shapes, rng, diff_modes, test_name=None):
     op_record("log1p", 1, numeric_dtypes, all_shapes, jtu.rand_small_positive(), []),
     op_record("logaddexp", 2, float_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
     op_record("logaddexp2", 2, float_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
+    op_record("polyval", 2, numeric_dtypes, nonempty_nonscalar_array_shapes, jtu.rand_default(), []),
     op_record("ravel", 1, default_dtypes, all_shapes, jtu.rand_default(), ["rev"]),
     op_record("remainder", 2, default_dtypes, all_shapes, jtu.rand_nonzero(), []),
     op_record("sqrt", 1, default_dtypes, all_shapes, jtu.rand_positive(), ["rev"]),
@@ -451,6 +453,28 @@ def args_maker():
     self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=True)
     self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_axis={}_baseshape=[{}]_dtypes=[{}]".format(
+          axis, ",".join(str(d) for d in base_shape),
+          ",".join(onp.dtype(dtype).name for dtype in dtypes)),
+       "axis": axis, "base_shape": base_shape, "dtypes": dtypes,
+       "rng": jtu.rand_default()}
+      for dtypes in CombosWithReplacement(default_dtypes, 2)
+      for base_shape in [(4,), (3, 4), (2, 3, 4)]
+      for axis in range(-len(base_shape)+1, len(base_shape))))
+  def testAppend(self, axis, base_shape, dtypes, rng):
+    wrapped_axis = axis % len(base_shape)
+    shapes = [base_shape[:wrapped_axis] + (size,) + base_shape[wrapped_axis+1:]
+              for size, _ in zip(itertools.cycle([3, 1, 4]), dtypes)]
+    onp_fun = lambda arr, values: onp.append(arr, values, axis=axis)
+    lnp_fun = lambda arr, values: lnp.append(arr, values, axis=axis)
+
+    def args_maker():
+      return [rng(shape, dtype) for shape, dtype in zip(shapes, dtypes)]
+
+    self._CheckAgainstNumpy(onp_fun, lnp_fun, args_maker, check_dtypes=True)
+    self._CompileAndCheck(lnp_fun, args_maker, check_dtypes=True)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape=[{}]_axis={}_repeats={}".format(
           jtu.format_shape_dtype_string(shape, dtype), axis, repeats),