anon-iclr-submission/benchname-bug-localization huggingface dataset:
Repository: google/jax
Issue: Regression bug in associative_scan

Hi,

There has been a regression in JAX code, which I believe to be related to this [commit](https://github.com/google/jax/commit/d3db7bd4be96e1fc16daa4571ade7fbfce13197f#diff-cc3f3a1265e2fee9115dc5a9f619d402e280fb48a7ed7d91261b42813d200f1b). 

I have created PR https://github.com/google/jax/pull/5165 with a reproducing failing test.

I will try and look further into what the problem could be, but if the person who made said commit could also have a look it would be great (@hawkinsp I believe this is you).

Adrien

Changed files:
tests/lax_vmap_test.py, jax/_src/lax/lax.py

diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
index 557a47acb..6183084f2 100644
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -4352,7 +4352,7 @@ def _gather_batching_rule(batched_args, batch_dims, *, dimension_numbers,
     counts = broadcasted_iota(start_indices.dtype, tuple(count_shape), 0)
     start_indices = concatenate([counts, start_indices], len(count_shape) - 1)
 
-    slice_sizes = (1,) + slice_sizes
+    slice_sizes = (_min(operand.shape[0], 1),) + slice_sizes
     collapsed_slice_dims = (0,) + tuple(np.add(1, dimension_numbers.collapsed_slice_dims))
     offset_dims = tuple(np.add(1, dimension_numbers.offset_dims))
     start_index_map = (0,) + tuple(np.add(1, dimension_numbers.start_index_map))
diff --git a/tests/lax_vmap_test.py b/tests/lax_vmap_test.py
index d38d4c9df..2c847eecf 100644
--- a/tests/lax_vmap_test.py
+++ b/tests/lax_vmap_test.py
@@ -74,7 +74,12 @@ class LaxVmapTest(jtu.JaxTestCase):
     args = [rng(shape, dtype) for shape, dtype in zip(batched_shapes, dtypes)]
     args_slice = args_slicer(args, bdims)
     ans = api.vmap(op, bdims)(*args)
-    expected = np.stack([op(*args_slice(i)) for i in range(bdim_size)])
+    if bdim_size == 0:
+      args = [rng(shape, dtype) for shape, dtype in zip(shapes, dtypes)]
+      out = op(*args)
+      expected = np.zeros((0,) + out.shape, out.dtype)
+    else:
+      expected = np.stack([op(*args_slice(i)) for i in range(bdim_size)])
     self.assertAllClose(ans, expected, rtol=rtol, atol=atol)
 
   @parameterized.named_parameters(itertools.chain.from_iterable(
@@ -642,6 +647,8 @@ class LaxVmapTest(jtu.JaxTestCase):
       for bdims in all_bdims(shape, idxs.shape)))
   def testGather(self, shape, dtype, idxs, dnums, slice_sizes, bdims):
     fun = partial(lax.gather, dimension_numbers=dnums, slice_sizes=slice_sizes)
+    self._CheckBatching(fun, 0, bdims, [shape, idxs.shape], [dtype, idxs.dtype],
+                        jtu.rand_default(self.rng()))
     self._CheckBatching(fun, 5, bdims, [shape, idxs.shape], [dtype, idxs.dtype],
                         jtu.rand_default(self.rng()))
Repository: google/jax
Issue: linalg.solve doesn't respect broadcasting semantics

I'm attempting a batched matrix solve,
```python
import jax.numpy as jp
import numpy as np

A = np.eye(2)[np.newaxis]
b = np.ones((3, 2))
print(A.shape, b.shape)
print(np.linalg.solve(A, b))
print(jp.linalg.solve(A, b))
```
but I'm getting the following:
```
(1, 2, 2) (3, 2)
[[1. 1.]
 [1. 1.]
 [1. 1.]]
/Users/skainswo/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
~/nu/skainswo/research/gan_with_the_wind/mle_normal.py in <module>
     31 print(A.shape, b.shape)
     32 print(np.linalg.solve(A, b))
---> 33 print(jp.linalg.solve(A, b))

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/numpy/linalg.py in solve(a, b)
    247   x = b if a_ndims == b_ndims else b[..., None]
    248 
--> 249   permutation = lax_linalg.lu_pivots_to_permutation(pivots, m)
    250   x = x[..., permutation, :]
    251 

~/.local/share/virtualenvs/research-OGGq2tNy/lib/python3.7/site-packages/jax/lax_linalg.py in lu_pivots_to_permutation(swaps, k)
    462     return swaps, permutation
    463 
--> 464   n, = np.shape(swaps)
    465   permutation = np.arange(k)
    466   _, permutation = lax.fori_loop(

ValueError: too many values to unpack (expected 1)
```

I'm on jax 0.1.36 and numpy 1.16.4.

Changed files:
jax/_src/lax/linalg.py, tests/linalg_test.py

diff --git a/jax/_src/lax/linalg.py b/jax/_src/lax/linalg.py
index 49c4b9d5e..ac7ba3f5f 100644
--- a/jax/_src/lax/linalg.py
+++ b/jax/_src/lax/linalg.py
@@ -253,6 +253,12 @@ def _check_solve_shapes(a, b):
 def _solve(a, b):
   _check_solve_shapes(a, b)
 
+  # Broadcast leading dimensions of b to the shape of a, as is required by
+  # custom_linear_solve.
+  out_shape = tuple(d_a if d_b == 1 else d_b
+                    for d_a, d_b in zip(a.shape[:-1] + (1,), b.shape))
+  b = jnp.broadcast_to(b, out_shape)
+
   # With custom_linear_solve, we can reuse the same factorization when
   # computing sensitivities. This is considerably faster.
   lu_, _, permutation = lu(lax.stop_gradient(a))
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 6dabfe6f0..4f35fef50 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -706,7 +706,7 @@ class NumpyLinalgTest(jtu.JaxTestCase):
           ((4, 4), (4,)),
           ((8, 8), (8, 4)),
           ((1, 2, 2), (3, 2)),
-          ((2, 1, 3, 3), (2, 4, 3, 4)),
+          ((2, 1, 3, 3), (1, 4, 3, 4)),
           ((1, 0, 0), (1, 0, 2)),
       ]
       for dtype in float_types + complex_types))
Repository: google/jax
Issue: fix higher-order rev-mode lax.index_take test failures 

With the new-and-more-thorough `check_grads`, we started getting some failures in higher-order tests of `lax.index_take`. See commented-out line in lax_test.py's testIndexTake.

Changed files:
tests/lax_test.py, jax/lax/lax.py

diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 56b96e0b2..5feaa7924 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -662,6 +662,7 @@ def scatter(operand, scatter_indices, updates, dimension_numbers):
 
 def index_take(src, idxs, axes):
   indices = concatenate([reshape(i, [i.shape[0], 1]) for i in idxs], 1)
+  indices = indices % onp.array([src.shape[ax] for ax in axes])
   slice_sizes = list(src.shape)
   for ax in axes:
     slice_sizes[ax] = 1
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 4a5f47cef..184882646 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2038,12 +2038,9 @@ class LaxAutodiffTest(jtu.JaxTestCase):
       ]
       for rng in [jtu.rand_default()]))
   def testIndexTakeGrad(self, shape, dtype, idxs, axes, rng):
-    idxs = tuple(rng(e.shape, e.dtype) for e in idxs)
     src = rng(shape, dtype)
     index_take = lambda src: lax.index_take(src, idxs, axes)
-    check_grads(index_take, (src,), 2, ["fwd"], 1e-2, 1e-2, 1)
-    # TODO(mattjj): fix rev mode failures here!
-    # check_grads(index_take, (src,), 2, ["fwd", "rev"], 1e-2, 1e-2, 1)
+    check_grads(index_take, (src,), 2, ["fwd", "rev"], 1e-2, 1e-2, 1)
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_idxs={}_dnums={}_slice_sizes={}".format(
Repository: google/jax
Issue: `inf` using multivariate_normal.logpdf with multiple batches

`multivariate_normal.logpdf` gives `inf`s when there are multiple batches:

```
>>> from jax import scipy
>>> import jax.numpy as np
>>> from jax.scipy.stats import multivariate_normal
>>> multivariate_normal.logpdf(np.array([[1. ,2.],
...                                      [2., 3.]]),
...                            np.array([[1., 2.],
...                                      [2., 3.]]),
...                            np.array([[[1., 0.2],
...                                       [0.2, 1.]],
...                                      [[1., 0.3],
...                                       [0.3, 1.]]]))
DeviceArray([inf, inf], dtype=float32)
```








Changed files:
tests/scipy_stats_test.py, jax/_src/scipy/stats/multivariate_normal.py

diff --git a/jax/_src/scipy/stats/multivariate_normal.py b/jax/_src/scipy/stats/multivariate_normal.py
index 213832160..9d7049488 100644
--- a/jax/_src/scipy/stats/multivariate_normal.py
+++ b/jax/_src/scipy/stats/multivariate_normal.py
@@ -44,7 +44,7 @@ def logpdf(x, mean, cov, allow_singular=None):
       L = lax.linalg.cholesky(cov)
       y = lax.linalg.triangular_solve(L, x - mean, lower=True, transpose_a=True)
       return (-1/2 * jnp.einsum('...i,...i->...', y, y) - n/2*np.log(2*np.pi)
-              - jnp.log(L.diagonal()).sum())
+              - jnp.log(L.diagonal(axis1=-1, axis2=-2)).sum(-1))
 
 @_wraps(osp_stats.multivariate_normal.pdf, update_doc=False)
 def pdf(x, mean, cov):
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
index eaa447819..c84ebe450 100644
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -20,6 +20,7 @@ from absl.testing import absltest, parameterized
 import numpy as np
 import scipy.stats as osp_stats
 
+from jax import api
 from jax import test_util as jtu
 from jax.scipy import stats as lsp_stats
 from jax.scipy.special import expit
@@ -470,6 +471,24 @@ class LaxBackedScipyStatsTests(jtu.JaxTestCase):
     self._CompileAndCheck(lsp_stats.multivariate_normal.logpdf, args_maker,
                           rtol=1e-4, atol=1e-4)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_ndim={}_nbatch={}_dtype={}".format(ndim, nbatch, dtype.__name__),
+       "ndim": ndim, "nbatch": nbatch, "dtype": dtype}
+      for ndim in [2, 3]
+      for nbatch in [1, 3, 5]
+      for dtype in jtu.dtypes.floating))
+  def testMultivariateNormalLogpdfBatch(self, ndim, nbatch, dtype):
+    # Regression test for #5570
+    rng = jtu.rand_default(self.rng())
+    x = rng((nbatch, ndim), dtype)
+    mean = 5 * rng((nbatch, ndim), dtype)
+    factor = rng((nbatch, ndim, 2 * ndim), dtype)
+    cov = factor @ factor.transpose(0, 2, 1)
+
+    result1 = lsp_stats.multivariate_normal.logpdf(x, mean, cov)
+    result2 = api.vmap(lsp_stats.multivariate_normal.logpdf)(x, mean, cov)
+    self.assertArraysEqual(result1, result2)
+
 
 if __name__ == "__main__":
     absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax
Issue: bincount of empty array raises

Running
`jnp.bincount(jnp.int32([]))`
raises
`ValueError: zero-size array to reduction operation max which has no identity`

In contrast, `np.bincount([])` produces `array([], dtype=int64)`

Changed files:
jax/_src/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/_src/numpy/lax_numpy.py b/jax/_src/numpy/lax_numpy.py
index 4332d6690..d22465799 100644
--- a/jax/_src/numpy/lax_numpy.py
+++ b/jax/_src/numpy/lax_numpy.py
@@ -1659,7 +1659,7 @@ def bincount(x, weights=None, minlength=0, *, length=None):
     x = core.concrete_or_error(asarray, x,
       "The error occured because of argument 'x' of jnp.bincount. "
       "To avoid this error, pass a static `length` argument.")
-    length = max(x) + 1
+    length = max(x, initial=-1) + 1
   length = _max(length, minlength)
   if ndim(x) != 1:
     raise ValueError("only 1-dimensional input supported.")
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index f42191ba9..a831c2ec1 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -4001,7 +4001,7 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
       "weights": weights,
       "minlength": minlength,
       "length": length}
-    for shape in [(5,), (10,)]
+    for shape in [(0,), (5,), (10,)]
     for dtype in int_dtypes
     for weights in [True, False]
     for minlength in [0, 20]
Repository: google/jax
Issue: jax.devices() behaves unintuitively

Hi! I noticed some unexpected (?) behaviour in the following code:
```python 
import jax
from jax.config import config
config.update("jax_backend_target", 'grpc://some endpoint:8470')
config.update("jax_xla_backend", 'tpu_driver')
print(jax.devices('cpu')) #prints tpu devices instead of cpu
``` 
I can get the expected behaviour if I do 

```python 
import jax
print(jax.devices('cpu')) #prints cpu devices 
from jax.config import config
config.update("jax_backend_target", 'grpc://11.111.11.1118470')
config.update("jax_xla_backend", 'tpu_driver')
print(jax.devices('cpu')) #now prints cpu devices
print(jax.devices('tpu')) #now prints tpu devices
``` 
I think this may be related the caching of `jax.lib.xla_bridge.get_backend()`. Not sure if this is expected behaviour or a bug. 

I noticed this because I was trying `jit` a few smaller functions on the host vm cpu during a larger TPU computation . I tried using the `backend=cpu` argument and `device_put`, but was unable to obtain the desired behaviour. In the end the only thing that seemed to work was to clear the cache of `get_backend()` reconfigure `jax.config` to cpu.

Changed files:
jax/lib/xla_bridge.py

diff --git a/jax/lib/xla_bridge.py b/jax/lib/xla_bridge.py
index aba3c6615..a292730c3 100644
--- a/jax/lib/xla_bridge.py
+++ b/jax/lib/xla_bridge.py
@@ -144,7 +144,9 @@ register_backend('xla', _get_local_backend)
 _tpu_backend = None
 
 def _get_tpu_driver_backend(platform):
-  del platform
+  if platform == "cpu":
+    return _get_local_backend("cpu")
+
   global _tpu_backend
   if _tpu_backend is None:
     backend_target = FLAGS.jax_backend_target
Repository: google/jax
Issue: Bug in jax.random.poisson for lam=0

Dear people from JAX,

Thank you for writing this software, I love it! I am using JAX for some optics simulations for which I need to use the poissonian random number generator to add noise to my simulated measurements. I noticed that `poisson(key, lam=0)` always returns -1, which makes no sense, and it does not match with the numpy implementation.

In JAX, I would write:

```python
import jax.numpy as jnp
import jax
key = jax.random.PRNGKey(0)
original = jnp.zeros(10,)
x = jax.random.poisson(key,original,(10,))
print(x)
# [-1 -1 -1 -1 -1 -1 -1 -1 -1 -1]
```
Numpy:

```python
import numpy as np
rng = np.random.default_rng()
s = rng.poisson(0, 10)
print(s)
# [0 0 0 0 0 0 0 0 0 0]
```

Looking at the source the problem seems to be in _poisson_knuth. An obvious solution would be to add a statement that just returns zero wherever lam==0,  but there's probably a more elegant solution?

Changed files:
jax/_src/random.py, tests/random_test.py

diff --git a/jax/_src/random.py b/jax/_src/random.py
index d6447cd8e..778080067 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -1173,11 +1173,12 @@ def _poisson(key, lam, shape, dtype):
   # λ -> ∞, so pick some arbitrary large value.
   lam_rejection = lax.select(use_knuth, lax.full_like(lam, 1e5), lam)
   max_iters = dtype.type(jnp.iinfo(dtype).max)  # insanely conservative
-  return lax.select(
-      use_knuth,
-      _poisson_knuth(key, lam_knuth, shape, dtype, max_iters),
-      _poisson_rejection(key, lam_rejection, shape, dtype, max_iters),
+  result = lax.select(
+    use_knuth,
+    _poisson_knuth(key, lam_knuth, shape, dtype, max_iters),
+    _poisson_rejection(key, lam_rejection, shape, dtype, max_iters),
   )
+  return lax.select(lam == 0, jnp.zeros_like(result), result)
 
 
 def poisson(key, lam, shape=(), dtype=dtypes.int_):
diff --git a/tests/random_test.py b/tests/random_test.py
index d1174c437..3b7dc7e12 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -538,6 +538,12 @@ class LaxRandomTest(jtu.JaxTestCase):
     x = random.poisson(key, np.array([2.0, 20.0]), shape=(3, 2))
     assert x.shape == (3, 2)
 
+  def testPoissonZeros(self):
+    key = random.PRNGKey(0)
+    lam = jnp.concatenate([jnp.zeros(10), 20 * jnp.ones(10)])
+    samples = random.poisson(key, lam, shape=(2, 20))
+    self.assertArraysEqual(samples[:, :10], jnp.zeros_like(samples[:, :10]))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_dtype={}".format(np.dtype(dtype).name), "dtype": dtype}
       for dtype in jtu.dtypes.floating))
Repository: google/jax
Issue: `AssertionError` with new `custom_jvp` for `jax.nn.relu`

On this [commit](https://github.com/google/jax/commit/ca23be63fbaed20192cda5f921afe177ac8dcf4d) (currently just a few commits behind head, none of which look relevant), I'm finding that
```python
import jax
import jax.numpy as np
import collections

class F(collections.namedtuple("F", ["a"])):
  def __call__(self, x):
    return jax.nn.relu(self.a) * x

@jax.jit
def g(f, x):
  return f(x)

jax.grad(F(2.0))(0.)  # works, returns DeviceArray(2., dtype=float32)
jax.grad(g, argnums=(1,))(F(2.0), 0.)  # AssertionError
```
hits an `AssertionError` on this line: https://github.com/google/jax/blob/1bb9aaa88c09a90a115c0304c05b0fc25932523b/jax/interpreters/partial_eval.py#L239

Commenting out the `@jax.jit` on `g` or replacing `jax.nn.relu(self.a)` with `np.maximum(self.a, 0)` both avoid the error.

Changed files:
jax/custom_derivatives.py, tests/api_test.py

diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
index bf8e9c5da..6e60bcd64 100644
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -252,6 +252,7 @@ def _flatten_jvp(in_tree, *args):
   yield primals_out + tangents_out, out_tree
 
 def _custom_jvp_call_bind(prim, fun, jvp, *args):
+  args = map(core.full_lower, args)
   top_trace = core.find_top_trace(args)
   level = (core.trace_state.trace_stack.next_level(True)
            if top_trace is None else top_trace.level)
@@ -490,6 +491,7 @@ def _flatten_bwd(in_tree, out_trees, *args):
   yield cts_in
 
 def _custom_vjp_call_bind(prim, fun, fwd, bwd, *args, out_trees):
+  args = map(core.full_lower, args)
   top_trace = core.find_top_trace(args)
   level = (core.trace_state.trace_stack.next_level(True)
            if top_trace is None else top_trace.level)
diff --git a/tests/api_test.py b/tests/api_test.py
index 365f4149e..ccdc5218c 100644
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2781,6 +2781,19 @@ class CustomVJPTest(jtu.JaxTestCase):
     expected = jax.grad(f, 0)(2., 0.1) + jax.grad(f, 0)(2., 0.2)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def test_lowering_out_of_traces(self):
+    # https://github.com/google/jax/issues/2578
+
+    class F(collections.namedtuple("F", ["a"])):
+      def __call__(self, x):
+        return jax.nn.relu(self.a) * x
+
+    @jax.jit
+    def g(f, x):
+      return f(x)
+
+    jax.grad(g, argnums=(1,))(F(2.0), 0.)  # doesn't crash
+
 
 class DeprecatedCustomTransformsTest(jtu.JaxTestCase):
Repository: google/jax
Issue: BatchNorm doesn't work with center or scale turned off

BatchNorm raises TypeError when using without center or scale or both. Seems pretty straightforward to fix (e.x. as a workaround I've set beta to None `beta = None if not center else beta[ed]` if center is False, the same for gamma but there might be some better solution).

Exception:
beta = beta[ed] # ed=(None, slice(None, None, None))
TypeError: tuple indices must be integers or slices, not tuple

Code:
```
from jax import np
from jax import random
from jax.experimental import stax


init_fn, apply_fn = stax.BatchNorm(axis=0, center=False, scale=False)
_, params = init_fn(random.PRNGKey(0), (-1, 1))

x = np.zeros((2, 1))
apply_fn(params, x)
```

Details:
jax==0.1.62
jaxlib==0.1.42

Changed files:
jax/experimental/stax.py, tests/stax_test.py

diff --git a/jax/experimental/stax.py b/jax/experimental/stax.py
index d2413181c..b5881a0f6 100644
--- a/jax/experimental/stax.py
+++ b/jax/experimental/stax.py
@@ -134,12 +134,10 @@ def BatchNorm(axis=(0, 1, 2), epsilon=1e-5, center=True, scale=True,
     # TODO(phawkins): np.expand_dims should accept an axis tuple.
     # (https://github.com/numpy/numpy/issues/12290)
     ed = tuple(None if i in axis else slice(None) for i in range(np.ndim(x)))
-    beta = beta[ed]
-    gamma = gamma[ed]
     z = normalize(x, axis, epsilon=epsilon)
-    if center and scale: return gamma * z + beta
-    if center: return z + beta
-    if scale: return gamma * z
+    if center and scale: return gamma[ed] * z + beta[ed]
+    if center: return z + beta[ed]
+    if scale: return gamma[ed] * z
     return z
   return init_fun, apply_fun
 
diff --git a/tests/stax_test.py b/tests/stax_test.py
index 0c009a593..27d9dc858 100644
--- a/tests/stax_test.py
+++ b/tests/stax_test.py
@@ -212,6 +212,19 @@ class StaxTest(jtu.JaxTestCase):
     assert out_shape == out.shape
     assert onp.allclose(onp.sum(onp.asarray(out), -1), 1.)
 
+  def testBatchNormNoScaleOrCenter(self):
+    key = random.PRNGKey(0)
+    axes = (0, 1, 2)
+    init_fun, apply_fun = stax.BatchNorm(axis=axes, center=False, scale=False)
+    input_shape = (4, 5, 6, 7)
+    inputs = random_inputs(onp.random.RandomState(0), input_shape)
+
+    out_shape, params = init_fun(key, input_shape)
+    out = apply_fun(params, inputs)
+    means = onp.mean(out, axis=(0, 1, 2))
+    std_devs = onp.std(out, axis=(0, 1, 2))
+    assert onp.allclose(means, onp.zeros_like(means), atol=1e-4)
+    assert onp.allclose(std_devs, onp.ones_like(std_devs), atol=1e-4)
 
   def testBatchNormShapeNHWC(self):
     key = random.PRNGKey(0)
Repository: google/jax
Issue: jaxlib build w/ cuda: File not found during compilation

I'm compiling `jaxlib` with CUDA 10.0 on Ubuntu 18.04. The build fails with the following error:

```
$ python3 build/build.py --enable_cuda --cuda_path /usr/local/cuda-10.0/ --cudnn_path /usr/local/cuda-10.0/ --enable_march_native

[...]

ERROR: /home/clem/.cache/bazel/_bazel_clem/ffaac3f7c6ad1cb26f04f1933452eef6/external/nccl_archive/BUILD.bazel:53:1: error while parsing .d file: /h
ome/clem/.cache/bazel/_bazel_clem/ffaac3f7c6ad1cb26f04f1933452eef6/execroot/__main__/bazel-out/k8-opt/bin/external/nccl_archive/_objs/device_lib/pr
od_i32_reduce_scatter.cu.d (No such file or directory)
nvcc fatal   : Could not open input file /tmp/tmpxft_00000004_00000000-6_prod_i32_reduce_scatter.cu.compute_35.cpp1.ii
Target //build:install_xla_in_source_tree failed to build
INFO: Elapsed time: 278.116s, Critical Path: 69.60s
INFO: 1281 processes: 1281 linux-sandbox.
FAILED: Build did NOT complete successfully
FAILED: Build did NOT complete successfully
Traceback (most recent call last):
  File "build/build.py", line 331, in <module>
    main()
  File "build/build.py", line 326, in main
    [":install_xla_in_source_tree", os.getcwd()])
  File "build/build.py", line 50, in shell
    output = subprocess.check_output(cmd)
  File "/usr/lib/python3.6/subprocess.py", line 356, in check_output
    **kwargs).stdout
  File "/usr/lib/python3.6/subprocess.py", line 438, in run
    output=stdout, stderr=stderr)
subprocess.CalledProcessError: Command '['./bazel-0.24.1-linux-x86_64', 'run', '--verbose_failures=true', '--config=opt', '--config=mkl_open_source
_only', '--config=cuda', ':install_xla_in_source_tree', '/home/clem/git/jax/build']' returned non-zero exit status 1.
```

Above this error message are only compiler warnings but no errors which could lead to some file not being created. Am I missing something? Or might there be a file name bug? Thanks a lot for your help!

---

I'm on a fresh Ubuntu 18.04.2 install with CUDA 10.0, cudnn and driver version 410.48.
[Full log](http://paste.ubuntu.com/p/tvXBHbr5gw/)

Changed files:
build/build.py

diff --git a/build/build.py b/build/build.py
index f2428415c..32fb7b38d 100755
--- a/build/build.py
+++ b/build/build.py
@@ -187,6 +187,9 @@ build --define=grpc_no_ares=true
 
 build:cuda --crosstool_top=@local_config_cuda//crosstool:toolchain
 build:cuda --define=using_cuda=true --define=using_cuda_nvcc=true
+
+build --spawn_strategy=standalone
+build --strategy=Genrule=standalone
 """
Repository: google/jax
Issue: invalid jaxpr from composition of soft_pmap, _papply, jit, and jit

The following:
```
@jit
def f(x):
  y = 2 * x
  @jit
  def g(z):
    h, axis_name =_papply(lambda x: x * y)
    return soft_pmap(h, axis_name)(z)
  return g(x)

f(onp.arange(4.).reshape((2, 2)))
```
results in the following invalid jaxpr:
```
{ lambda b ;  ; a.
  let c = reshape[ new_sizes=(1, 2, 2)
                   dimensions=None
                   old_sizes=(2, 2) ] a
      d = mul e 2.0
      f = xla_pmap[ axis_size=1
                    axis_name=<axis 0x7fa985eaa890> ] b c
        { lambda  ; b ; c a.
          let d = gather[ operand_shape=(2, 2)
                          slice_sizes=(1, 2)
                          dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2), collapsed_slice_dims=(), start_index_map=(0, 1)) ] b c
              e = reshape[ new_sizes=(2, 2)
                           dimensions=None
                           old_sizes=(2, 1, 2) ] d
              f = mul a e
          in f } [  ; d ]
      g = reshape[ new_sizes=(2, 2)
                   dimensions=None
                   old_sizes=(1L, 2, 2) ] f
  in g }
```

Changed files:
jax/interpreters/partial_eval.py

diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
index 0da6b8cb9..7195c57f6 100644
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -116,8 +116,7 @@ class JaxprTrace(Trace):
     out_pv_const, consts = call_primitive.bind(fun, *in_consts, **params)
     out_pv, jaxpr, env = aux()
     const_tracers = map(self.new_instantiated_const, consts)
-    env_tracers = map(self.full_raise, env)
-    bound_subjaxpr = (jaxpr, const_tracers, env_tracers)
+    bound_subjaxpr = (jaxpr, const_tracers, map(self.full_raise, env))
     eqn = JaxprEqn(tracers, None, call_primitive, (bound_subjaxpr,),
                    False, False, params)
     return JaxprTracer(self, PartialVal((out_pv, out_pv_const)), eqn)
@@ -130,12 +129,11 @@ class JaxprTrace(Trace):
     out_pv_reduced, jaxpr, env = aux()
     out_pv = add_axis_to_pv(params['axis_size'], out_pv_reduced)
     const_tracers = map(self.new_instantiated_const, consts)
-    env_tracers = map(self.full_raise, env)
     jaxpr_converted = jaxpr.copy()
     jaxpr_converted.constvars = []
     jaxpr_converted.invars = list(it.chain(jaxpr.constvars, jaxpr.invars))
     invars = tuple(it.chain(const_tracers, tracers))
-    bound_subjaxpr = (jaxpr_converted, (), env)
+    bound_subjaxpr = (jaxpr_converted, (), map(self.full_raise, env))
     eqn = JaxprEqn(invars, None, map_primitive, (bound_subjaxpr,),
                    False, False, params)
     return JaxprTracer(self, PartialVal((out_pv, out_const)), eqn)
@@ -465,6 +463,7 @@ def tracers_to_jaxpr(in_tracers, out_tracer):
     if isinstance(recipe, JaxprEqn):
       eqns.append(eqn_tracer_to_var(var, [var(t)], recipe))
     elif isinstance(recipe, LambdaBinding):
+      assert any(t is in_tracer for in_tracer in in_tracers)
       assert in_tracers, "Lambda binding with no args"
     elif isinstance(recipe, FreeVar):
       env[var(t)] = recipe.val
Repository: google/jax
Issue: Inconsistent convention for complex derivatives

The convention for the handling of complex derivatives in JAX seems to be mixed. Consider the following two implementations of the abolute value function.
```py
import jax.numpy as np
from jax import grad

def f(x):
    return np.abs(x)

def g(x):
    return np.sqrt(np.real(x)**2 + np.imag(x)**2)

df = grad(f)
dg = grad(g)

print('gradient by df: ',df(1.0+0.5j))
print('gradient by dg: ',dg(1.0+0.5j))
```
The respective results differ by complex conjugation:
```
gradient by df:  (0.8944272-0.4472136j)
gradient by dg:  (0.8944272+0.4472136j)
```

Changed files:
tests/lax_test.py, jax/lax/lax.py

diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 675e6200f..26fff8a98 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1608,12 +1608,12 @@ real_p = unop(_complex_basetype, _complex, 'real')
 ad.deflinear(real_p, lambda t: [complex(t, onp.zeros((), _dtype(t)))])
 
 imag_p = unop(_complex_basetype, _complex, 'imag')
-ad.deflinear(imag_p, lambda t: [complex(onp.zeros((), _dtype(t)), t)])
+ad.defjvp(imag_p, lambda g, _: real(mul(_const(g, -1j), g)))
 
 _complex_dtype = lambda dtype, *args: (onp.zeros((), dtype) + onp.zeros((), onp.complex64)).dtype
 complex_p = binop(_complex_dtype, [_complex_elem_types, _complex_elem_types],
                   'complex')
-ad.deflinear(complex_p, lambda t: [real(t), imag(t)])
+ad.deflinear(complex_p, lambda t: [real(t), imag(neg(t))])
 
 conj_p = unop(_complex_dtype, _float | _complex, 'conj')
 
diff --git a/tests/lax_test.py b/tests/lax_test.py
index 3da63b8a9..aafa32193 100644
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1523,10 +1523,10 @@ LAX_GRAD_OPS = [
 
     grad_test_spec(lax.real, nargs=1, order=2, rng=jtu.rand_default(),
                    dtypes=[onp.complex64]),
-    # grad_test_spec(lax.imag, nargs=1, order=2, rng=jtu.rand_default(),
-    #                dtypes=[onp.complex64]),  # TODO(mattjj): enable
-    # grad_test_spec(lax.complex, nargs=2, order=2, rng=jtu.rand_default(),
-    #                dtypes=[onp.float32]),  # TODO(mattjj): enable
+    grad_test_spec(lax.imag, nargs=1, order=2, rng=jtu.rand_default(),
+                   dtypes=[onp.complex64]),
+    grad_test_spec(lax.complex, nargs=2, order=2, rng=jtu.rand_default(),
+                   dtypes=[onp.float32]),
     grad_test_spec(lax.conj, nargs=1, order=2, rng=jtu.rand_default(),
                    dtypes=[onp.float32, onp.complex64]),
     grad_test_spec(lax.abs, nargs=1, order=2, rng=jtu.rand_positive(),
Repository: google/jax
Issue: Key word arguments breaks pmap

This code
```python
def f(index1, index2):
  return np.arange(36).reshape(6, 6)[index1, index2]
g = jax.jit(jax.pmap(f))
g(index1=onp.asarray([1]), index2=onp.asarray([2]))
```

Causes this error.
```
google3/third_party/py/jax/api.py in _pmap_axis_size(args)
    724     return next(_axis_size(leaf) for arg in args for leaf in tree_leaves(arg))
    725   except StopIteration:
--> 726     raise ValueError("pmap requires a leading axis to map over.")
    727 
    728 def _axis_size(x):

ValueError: pmap requires a leading axis to map over.
```

Whereas the exact same code without the kwargs works correctly.

```python
 
def f(index1, index2):
  return np.arange(36).reshape(6, 6)[index1, index2]
g = jax.jit(jax.pmap(f))
g(onp.asarray([1]), onp.asarray([2]))
```
Which nicely prints
```
DeviceArray([8], dtype=int32)
```
as expected.

Changed files:
tests/batching_test.py, jax/api.py

diff --git a/jax/api.py b/jax/api.py
index d8e96e258..8609a144d 100644
--- a/jax/api.py
+++ b/jax/api.py
@@ -711,7 +711,7 @@ def pmap(fun, axis_name=None):
 
   @wraps(fun)
   def f_pmapped(*args, **kwargs):
-    axis_size = _pmap_axis_size(args)
+    axis_size = _pmap_axis_size((args, kwargs))
     f = lu.wrap_init(fun)
     args_flat, in_tree = tree_flatten((args, kwargs))
     flat_fun, out_tree = flatten_fun_leafout(f, in_tree)
diff --git a/tests/batching_test.py b/tests/batching_test.py
index cf4aadeee..898ff5e52 100644
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -955,6 +955,14 @@ class BatchingTest(jtu.JaxTestCase):
     result = vmap(f, (None, 0, None))(onp.zeros((10,)), onp.arange(10,), 1.)
     self.assertAllClose(result, onp.eye(10), check_dtypes=False)
 
+  def testIssue1170(self):
+    def f(index1, index2):
+      return np.arange(36).reshape(6, 6)[index1, index2]
+    g = jax.jit(jax.pmap(f))
+    ans = g(index1=onp.asarray([1]), index2=onp.asarray([2]))
+    expected = g(onp.asarray([1]), onp.asarray([2]))
+    self.assertAllClose(ans, expected, check_dtypes=True)
+
 
 if __name__ == '__main__':
   absltest.main()
Repository: google/jax
Issue: np.arange ignores dtype

arange ignores the passed dtype and is always `int*`
```python
a = np.arange(4, dtype=np.complex64)
print(a.dtype)
assert a.dtype == np.complex64
```
Results in:
```
int64
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
<ipython-input-12-5b101f76f4c3> in <module>()
      1 a = np.arange(4, dtype=np.complex64)
      2 print(a.dtype)
----> 3 assert a.dtype is np.complex64

AssertionError: 
```

Changed files:
jax/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index 4efb257f1..2c433da8d 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -1379,7 +1379,7 @@ def identity(n, dtype=None):
 
 @_wraps(onp.arange)
 def arange(*args, **kwargs):
-  dtype = kwargs.pop("dtype", None)
+  dtype = kwargs.get("dtype", None)
   if not args:
     raise TypeError("Required argument 'start' (pos 1) not found")  # same as numpy error
 
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 67273f5a4..bdc834deb 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1493,7 +1493,6 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     expected = onp.reshape(a, (3, 2), order='F')
     self.assertAllClose(ans, expected, check_dtypes=True)
 
-
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_op={}_dtype={}".format(
           op, {bool: "bool", int: "int", float: "float"}[dtype]),
@@ -1538,6 +1537,10 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     self.assertFalse(type(lnp.arange(77)) == type(onp.arange(77)))
     self.assertTrue(type(lnp.arange(77)) == type(lax.iota(onp.int32, 77)))
 
+  def testIssue830(self):
+    a = lnp.arange(4, dtype=lnp.complex64)
+    self.assertEqual(a.dtype, lnp.complex64)
+
   def testIssue728(self):
     assert lnp.allclose(lnp.eye(5000), onp.eye(5000))
     self.assertEqual(0, onp.sum(lnp.eye(1050) - onp.eye(1050)))
Repository: google/jax
Issue: np.stack raises incorrect exception when given a traced array 

```python
@jax.jit
def foo(x):
  return np.stack(x)

foo(np.zeros(2))  # this fails
```
Raises an exception because of a `if not arrays` check which assumes `arrays` is an object without a `__bool__` method so that the size of arrays would be used instead. Using `len(arrays) == 0` might be sufficient to fix this. This issue might be present in a few other places and might warrant a test.

Error message:
```python
TypeError: Abstract value passed to `bool`, which requires a concrete value. The function to be transformed can't be traced at the required level of abstraction. If using `jit`, try using `static_argnums` or applying `jit` to smaller subfunctions instead.
```

Changed files:
jax/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index 6006f93b9..0c14536ef 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -1352,7 +1352,7 @@ def pad(array, pad_width, mode='constant', constant_values=0):
 
 @_wraps(onp.stack)
 def stack(arrays, axis=0):
-  if not arrays:
+  if not len(arrays):
     raise ValueError("Need at least one array to stack.")
   shape0 = shape(arrays[0])
   axis = _canonicalize_axis(axis, len(shape0) + 1)
@@ -1377,7 +1377,7 @@ def tile(a, reps):
 
 @_wraps(onp.concatenate)
 def concatenate(arrays, axis=0):
-  if not arrays:
+  if not len(arrays):
     raise ValueError("Need at least one array to concatenate.")
   if ndim(arrays[0]) == 0:
     raise ValueError("Zero-dimensional arrays cannot be concatenated.")
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index f43167b91..b03f4fe3c 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1864,6 +1864,18 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
     finally:
       FLAGS.jax_numpy_rank_promotion = prev_flag
 
+  def testStackArrayArgument(self):
+    # tests https://github.com/google/jax/issues/1271
+    @api.jit
+    def foo(x):
+      return lnp.stack(x)
+    foo(onp.zeros(2))  # doesn't crash
+
+    @api.jit
+    def foo(x):
+      return lnp.concatenate(x)
+    foo(onp.zeros((2, 2)))  # doesn't crash
+
 
 # Most grad tests are at the lax level (see lax_test.py), but we add some here
 # as needed for e.g. particular compound ops of interest.
Repository: google/jax
Issue: forward-mode eigh derivatives return NaN's

First off, thanks for developing JAX, it's awesome!

I recently noticed that when one computes higher order derivatives of expressions involving eigenvectors obtained from `jax.numpy.eigh` or `jax.scipy.eigh`, NaNs are returned. However, it only appears to be a problem with forward-mode autodiff. This is unfortunate, since forward-mode tends to be more efficient for higher order derivatives.

A minimal example:
```python
import jax
import jax.numpy as np

# Create a random symmetric matrix
seed = jax.random.PRNGKey(0)
tmp = jax.random.uniform(seed, (2,2))
a = np.dot(tmp, tmp.T)

def test(inp):
    val, vec = np.linalg.eigh(inp)
    return np.dot(np.dot(vec, inp), vec.T)

def test_deriv(func, inp):
    grad_func = jax.jacfwd(func)
    hess_func = jax.jacfwd(grad_func)
    cube_func = jax.jacfwd(hess_func) # This derivative returns NaN, but jax.jacrev works!
    print(grad_func(inp))
    print(hess_func(inp))
    print(cube_func(inp))

test_deriv(test, a)
```

Additional notes:
- Expressions involving eigenvalues do not appear to have this issue
- Any array contractions involving eigenvectors and other arrays (`dot`, `einsum`, etc) also exhibit similar behavior: `jacfwd` fails, while `jacrev` works.
- I'm using CPU
- JAX version 0.1.46

Thanks for your time!

Changed files:
tests/linalg_test.py, jax/lax/lax.py

diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 713f94f55..a0318eed9 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1720,7 +1720,7 @@ def _div_transpose_rule(cotangent, x, y):
 div_p = standard_binop([_num, _num], 'div')
 ad.defjvp(div_p,
           lambda g, x, y: div(_brcast(g, y), y),
-          lambda g, x, y: div(mul(neg(_brcast(g, x)), x), pow(y, _two(y))))
+          lambda g, x, y: div(mul(neg(_brcast(g, x)), x), square(y)))
 ad.primitive_transposes[div_p] = _div_transpose_rule
 
 rem_p = standard_binop([_num, _num], 'rem')
diff --git a/tests/linalg_test.py b/tests/linalg_test.py
index 18f3688df..5762f4120 100644
--- a/tests/linalg_test.py
+++ b/tests/linalg_test.py
@@ -538,6 +538,20 @@ class NumpyLinalgTest(jtu.JaxTestCase):
     jac0 = jax.jacobian(np.linalg.solve, argnums=0)(A[0], b[0])
     jac1 = jax.jacobian(np.linalg.solve, argnums=1)(A[0], b[0])
 
+  def testIssue1383(self):
+    seed = jax.random.PRNGKey(0)
+    tmp = jax.random.uniform(seed, (2,2))
+    a = np.dot(tmp, tmp.T)
+
+    def f(inp):
+      val, vec = np.linalg.eigh(inp)
+      return np.dot(np.dot(vec, inp), vec.T)
+
+    grad_func = jax.jacfwd(f)
+    hess_func = jax.jacfwd(grad_func)
+    cube_func = jax.jacfwd(hess_func)
+    self.assertFalse(onp.any(onp.isnan(cube_func(a))))
+
 
 class ScipyLinalgTest(jtu.JaxTestCase):
Repository: google/jax
Issue: np.take_along_axis gives wrong gradient

The following repro script shows the issue:
```
import jax.numpy as np
from jax import value_and_grad

idx = np.repeat(np.arange(3), 10).reshape((30, 1))

def f(x):
    y = x * np.arange(3.).reshape((1, 3))
    return np.take_along_axis(y, idx, -1).sum()

def g(x):
    y = x * np.arange(3.).reshape((1, 3))
    y = np.broadcast_to(y, (30, 3))
    return np.take_along_axis(y, idx, -1).sum()

print(value_and_grad(f)(1.))  # get 30, 0
print(value_and_grad(g)(1.))  # get 30, 30
```

I think `(30, 30)` is the correct answer.

Changed files:
jax/numpy/lax_numpy.py, tests/lax_numpy_test.py

diff --git a/jax/numpy/lax_numpy.py b/jax/numpy/lax_numpy.py
index 435380021..ff0a42dba 100644
--- a/jax/numpy/lax_numpy.py
+++ b/jax/numpy/lax_numpy.py
@@ -2405,7 +2405,7 @@ def _take_along_axis(arr, indices, axis):
   if axis is None:
     if ndim(indices) != 1:
       msg = "take_along_axis indices must be 1D if axis=None, got shape {}"
-      raise ValueError(msg.format(shape(indices)))
+      raise ValueError(msg.format(indices.shape))
     return take_along_axis(arr.ravel(), indices, 0)
   rank = ndim(arr)
   if rank != ndim(indices):
@@ -2413,11 +2413,19 @@ def _take_along_axis(arr, indices, axis):
     raise ValueError(msg.format(ndim(indices), ndim(arr)))
   axis = _canonicalize_axis(axis, rank)
 
-  arr_shape = list(shape(arr))
-  axis_size = arr_shape[axis]
-  arr_shape[axis] = 1
-  idx_shape = shape(indices)
-  out_shape = lax.broadcast_shapes(idx_shape, tuple(arr_shape))
+  def replace(tup, val):
+    lst = list(tup)
+    lst[axis] = val
+    return tuple(lst)
+
+  bcast_shape = lax.broadcast_shapes(replace(arr.shape, 1), replace(indices.shape, 1))
+  indices = broadcast_to(indices, replace(bcast_shape, indices.shape[axis]))
+  arr     = broadcast_to(arr,     replace(bcast_shape, arr.shape[axis]))
+
+  axis_size = arr.shape[axis]
+  arr_shape = replace(arr.shape, 1)
+  idx_shape = indices.shape
+  out_shape = lax.broadcast_shapes(idx_shape, arr_shape)
 
   index_dims = [i for i, idx in enumerate(idx_shape) if i == axis or idx != 1]
 
@@ -3144,4 +3152,4 @@ def _unstack(x):
   if x.ndim == 0:
     raise ValueError("Argument to _unstack must be non-scalar")
   return [lax.index_in_dim(x, i, keepdims=False) for i in range(x.shape[0])]
-setattr(DeviceArray, "_unstack", _unstack)
\\ No newline at end of file
+setattr(DeviceArray, "_unstack", _unstack)
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index e97d7eec7..56f75cd24 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1985,6 +1985,16 @@ class NumpyGradTests(jtu.JaxTestCase):
   def testOpGradSpecialValue(self, op, special_value):
     check_grads(op, (special_value,), 2, ["fwd", "rev"])
 
+  def testTakeAlongAxisIssue1521(self):
+    # https://github.com/google/jax/issues/1521
+    idx = lnp.repeat(lnp.arange(3), 10).reshape((30, 1))
+
+    def f(x):
+      y = x * lnp.arange(3.).reshape((1, 3))
+      return lnp.take_along_axis(y, idx, -1).sum()
+
+    check_grads(f, (1.,), order=1)
+
 
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax
Issue: np.broadcast_to regression in the latest JAX version

The new version 0.1.47 introduces a regression where `np.broadcast_to(1, (3, 2))` throws the error `AttributeError: 'int' object has no attribute 'ndim'`.

cc @mattjj @neerajprad

Changed files:
tests/lax_numpy_test.py, jax/lax/lax.py

diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index fe895aa1e..0a648671b 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -599,7 +599,7 @@ def broadcast(operand, sizes):
   return broadcast_p.bind(operand, sizes=tuple(sizes))
 
 def broadcast_in_dim(operand, shape, broadcast_dimensions):
-  if operand.ndim == len(shape) and not len(broadcast_dimensions):
+  if onp.ndim(operand) == len(shape) and not len(broadcast_dimensions):
     return operand
   if any(x < 0 or x >= len(shape) for x in broadcast_dimensions):
     msg = ("broadcast dimensions must be >= 0 and < ndim(shape), got {} for "
diff --git a/tests/lax_numpy_test.py b/tests/lax_numpy_test.py
index 305bfe6ae..c7974658f 100644
--- a/tests/lax_numpy_test.py
+++ b/tests/lax_numpy_test.py
@@ -1952,6 +1952,10 @@ class LaxBackedNumpyTests(jtu.JaxTestCase):
         ValueError, "Incompatible shapes for broadcasting: .*",
         lambda: lnp.broadcast_to(onp.ones((2, 3)), (1, 3)))
 
+  def testBroadcastToIntIssue1548(self):
+    self.assertAllClose(lnp.broadcast_to(1, (3, 2)), onp.ones((3, 2)),
+                        check_dtypes=False)
+
 # Most grad tests are at the lax level (see lax_test.py), but we add some here
 # as needed for e.g. particular compound ops of interest.
Repository: google/jax
Issue: Add pxla.ShardedDeviceArray to jaxval_zeros_likers

This is a one-line fix at https://github.com/google/jax/blob/master/jax/lax/lax.py#L1503

Just adding this as an issue as a reminder for @mattjj 

Changed files:
jax/lax/lax.py

diff --git a/jax/lax/lax.py b/jax/lax/lax.py
index 2e05b0b2b..12f5ba849 100644
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -1498,9 +1498,10 @@ def zeros_like_array(x):
   return full_like(x, 0)
 
 for t in itertools.chain(dtypes.python_scalar_dtypes.keys(), array_types,
-                         [xla.DeviceArray]):
+                         [xla.DeviceArray, pxla.ShardedDeviceArray]):
   ad_util.jaxval_adders[t] = add
 ad_util.jaxval_zeros_likers[xla.DeviceArray] = zeros_like_array
+ad_util.jaxval_zeros_likers[pxla.ShardedDeviceArray] = zeros_like_array
 
 
 ### primitives
Repository: google/jax
Issue: complex precision not respected in jax.random.normal

I am trying to generate random complex data of different precision and noticed the following:

```python
from jax.config import config
config.update("jax_enable_x64", True)
import jax

print(jax.random.normal(jax.random.PRNGKey(1), dtype=jax.numpy.complex128).dtype)
# complex128

print(jax.random.normal(jax.random.PRNGKey(1), dtype=jax.numpy.complex64).dtype)
# complex128
```
even though the second one supposed to generate `complex64`

With the corresponding real types (float32, float64) this works fine.

Changed files:
jax/_src/random.py, tests/random_test.py

diff --git a/jax/_src/random.py b/jax/_src/random.py
index 778080067..500c5b144 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -634,7 +634,7 @@ def _normal(key, shape, dtype) -> jnp.ndarray:
     dtype = dtypes.dtype_real(dtype)
     _re = _normal_real(key_re, shape, dtype)
     _im = _normal_real(key_im, shape, dtype)
-    return 1 / sqrt2 * (_re + 1j * _im)
+    return (_re + 1j * _im) / sqrt2
   else:
     return _normal_real(key, shape, dtype) # type: ignore
 
diff --git a/tests/random_test.py b/tests/random_test.py
index cd74c8dd6..5acd5b5c9 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -223,6 +223,7 @@ class LaxRandomTest(jtu.JaxTestCase):
     for samples in [uncompiled_samples, compiled_samples]:
       self._CheckKolmogorovSmirnovCDF(jnp.real(samples), scipy.stats.norm(scale=1/np.sqrt(2)).cdf)
       self._CheckKolmogorovSmirnovCDF(jnp.imag(samples), scipy.stats.norm(scale=1/np.sqrt(2)).cdf)
+      self.assertEqual(dtype, samples.dtype)
 
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_dtype={}".format(np.dtype(dtype).name), "dtype": dtype}
Repository: google/jax
Issue: jax.random.poisson hangs if lam is NaN

The following causes Python to hang:

```python
import jax
import numpy as np

key = jax.random.PRNGKey(0)
jax.random.poisson(key=key, lam=jnp.array([np.nan]), shape=(1,))
# hangs
```

My expectation is that when lam=np.nan, the output will be np.nan.


Changed files:
jax/_src/random.py, tests/random_test.py

diff --git a/jax/_src/random.py b/jax/_src/random.py
index 6b51fb208..a9d894a3f 100644
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -53,6 +53,9 @@ UINT_DTYPES = prng.UINT_DTYPES
 
 ### utilities
 
+def _isnan(x):
+  return lax.ne(x, x)
+
 def _check_prng_key(key):
   # TODO(frostig): remove once we always enable_custom_prng
   if type(key) is prng.PRNGKeyArray:
@@ -1096,7 +1099,7 @@ def _poisson(key, lam, shape, dtype):
   # https://github.com/numpy/numpy/blob/v1.18.3/numpy/random/src/distributions/distributions.c#L574
   # For lambda < 10, we use the Knuth algorithm; otherwise, we use transformed
   # rejection sampling.
-  use_knuth = lam < 10
+  use_knuth = _isnan(lam) | (lam < 10)
   lam_knuth = lax.select(use_knuth, lam, lax.full_like(lam, 0.0))
   # The acceptance probability for rejection sampling maxes out at 89% as
   # λ -> ∞, so pick some arbitrary large value.
diff --git a/tests/random_test.py b/tests/random_test.py
index c58589624..bbf7a9e3c 100644
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -791,6 +791,12 @@ class LaxRandomTest(jtu.JaxTestCase):
     samples = random.poisson(key, lam, shape=(2, 20))
     self.assertArraysEqual(samples[:, :10], jnp.zeros_like(samples[:, :10]))
 
+  def testPoissonCornerCases(self):
+    key = self.seed_prng(0)
+    lam = jnp.array([-1, 0, jnp.nan])
+    samples = random.poisson(key, lam, shape=(3,))
+    self.assertArraysEqual(samples, jnp.array([-1, 0, -1]))
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_dtype={}".format(np.dtype(dtype).name), "dtype": dtype}
       for dtype in jtu.dtypes.floating))