internlm/SWE-Fixer-Train-110K huggingface dataset:
Repository: google/jax

Issue:
float16 support
Add support for `np.float16`.


Discussion:

Resolving patch:
diff --git a/jax/abstract_arrays.py b/jax/abstract_arrays.py
--- a/jax/abstract_arrays.py
+++ b/jax/abstract_arrays.py
@@ -161,9 +161,11 @@ def zeros_like_array(x):
   dtype = xla_bridge.canonicalize_dtype(onp.result_type(x))
   return onp.broadcast_to(onp.array(0, dtype), onp.shape(x))
 
-array_types = [onp.ndarray, onp.float64, onp.float32, onp.complex64,
-               onp.complex128, onp.int64, onp.int32, onp.bool_, onp.uint64,
-               onp.uint32, complex, float, int, bool]
+array_types = [onp.ndarray, onp.float64, onp.float32, onp.float16,
+               onp.complex64, onp.complex128,
+               onp.int64, onp.int32, onp.int16, onp.int8,
+               onp.bool_, onp.uint64, onp.uint32, onp.uint16, onp.uint8,
+               complex, float, int, bool]
 
 for t in array_types:
   core.pytype_aval_mappings[t] = ConcreteArray


Test patch:
Repository: google/jax

Issue:
Reverse-mode differentiation rule for 'sub' not implemented
I get an error with the gradients of multivariate_normal.logpdf wrt the covariance (input and mean are fine). For example, this

```python
import jax.numpy as np
from jax import grad
from jax.scipy.stats import multivariate_normal

f = lambda x: multivariate_normal.logpdf(np.zeros(1), np.zeros(1), np.ones((1, 1))*x)
grad(f)(1.)  # fails
```

gives the error:

```
Traceback (most recent call last):
  File "[...]/jax/jax/interpreters/ad.py", line 149, in get_primitive_transpose
    return primitive_transposes[p]
KeyError: sub

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "[...]", line 43, in <module>
    print(grad(f)(1.))  # fails
  File "[...]/jax/jax/api.py", line 149, in grad_f
    ans, g = value_and_grad_f(*args, **kwargs)
  File "[...]/jax/jax/api.py", line 186, in value_and_grad_f
    g = vjp_py(onp.ones((), onp.result_type(ans)))
  File "[...]/jax/jax/api_util.py", line 56, in apply_jaxtree_fun
    ans = fun(*args)
  File "[...]/jax/jax/api.py", line 457, in out_vjp_packed
    return out_vjp(cotangent_in)
  File "[...]/jax/jax/interpreters/ad.py", line 78, in vjp_
    _, arg_cts = backward_pass(jaxpr, consts, (), dummy_primal_and_ct)
  File "[...]/jax/jax/interpreters/ad.py", line 135, in backward_pass
    cts_out = get_primitive_transpose(eqn.primitive)(ct_in, *invals, **eqn.params)
  File "[...]/jax/jax/interpreters/ad.py", line 152, in get_primitive_transpose
    "Reverse-mode differentiation rule for '{}' not implemented".format(p))
NotImplementedError: Reverse-mode differentiation rule for 'sub' not implemented
```

Also, I notice that the multivariate_normal density is implemented with matrix inverses and determinant rather than the [usual](https://github.com/scipy/scipy/blob/d681fc87f42bbc237d5b7e98fbfbfb5e3a51ae13/scipy/stats/_multivariate.py#L495) more efficient formulation in terms of the square root cov. Is there a reason for this?


Discussion:

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1016,10 +1016,16 @@ def _add_transpose(t, x, y):
 ad.defjvp(add_p, lambda g, x, y: _brcast(g, y), lambda g, x, y: _brcast(g, x))
 ad.primitive_transposes[add_p] = _add_transpose
 
+
+def _sub_transpose(t, x, y):
+  assert x is None and y is None  # computation must be linear, not affine
+  return [t, neg(t)]
+
 sub_p = standard_binop([_num, _num], 'sub')
 ad.defjvp(sub_p,
           lambda g, x, y: _brcast(g, y),
           lambda g, x, y: _brcast(neg(g), x))
+ad.primitive_transposes[sub_p] = _sub_transpose
 
 mul_p = standard_binop([_num, _num], 'mul')
 ad.defbilinear_broadcasting(_brcast, mul_p, mul, mul)  # TODO


Test patch:
Repository: google/jax

Issue:
Error differentiating QR decomposition
This appears to be a bug but unfortunately I don't have time to delve into it.

Below is a simple script to do gradient descent on a function involving a QR decomposition. It fails when I execute `training_gradient_fun` with `TypeError: No abstraction handler for type: <class 'jax.ad_util.Zero'>`.

Ubuntu 16.04
python 3.6.6
jax 0.1.20
jaxlib 0.1.9

```
import jax
import jax.numpy as np
import numpy as onp

def loss(weights, inputs):
    q, r = np.linalg.qr(weights)
    a = np.matmul(q, weights)
    return np.mean(a ** 2)

training_gradient_fun = jax.jit(jax.grad(loss))

inputs = onp.random.rand(5, 5)
weights = onp.random.rand(5, 5)

for i in range(100):
    weights -= 0.1 * training_gradient_fun(weights, inputs)
```


Discussion:

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1363,7 +1363,7 @@ def _add_transpose(t, x, y):
 
 def _sub_transpose(t, x, y):
   assert x is None and y is None  # computation must be linear, not affine
-  return [t, neg(t)]
+  return [t, neg(t) if t is not ad_util.zero else ad_util.zero]
 
 sub_p = standard_binop([_num, _num], 'sub')
 ad.defjvp(sub_p,


Test patch:
Repository: google/jax

Issue:
`stax.BatchNorm` incorrect normalization
jax.experimental.stax on current master, line 140 ( `BatchNorm` function):
```python
z = (x - mean) / (var + epsilon)**2
```
The square in the denominator should probably be a square root.


Discussion:
Thanks so much for catching this!

Resolving patch:
diff --git a/jax/experimental/stax.py b/jax/experimental/stax.py
--- a/jax/experimental/stax.py
+++ b/jax/experimental/stax.py
@@ -137,7 +137,7 @@ def init_fun(input_shape):
   def apply_fun(params, x, rng=None):
     beta, gamma = params
     mean, var = np.mean(x, axis, keepdims=True), fastvar(x, axis, keepdims=True)
-    z = (x - mean) / (var + epsilon)**2
+    z = (x - mean) / np.sqrt(var + epsilon)
     if center and scale: return gamma * z + beta
     if center: return z + beta
     if scale: return gamma * z


Test patch:
Repository: google/jax

Issue:
Document `lax.stop_gradient`
Is there a function similar to tensorflow's [tf.stop_gradient](https://www.tensorflow.org/api_docs/python/tf/stop_gradient)?

How can I do this in jax??


Discussion:
It's not documented, but I think you're looking for `jax.lax.stop_gradient`.

We should document it!

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -1114,6 +1114,25 @@ def broadcasted_eye(dtype, shape, axes):
 
 
 def stop_gradient(x):
+  """Stops gradient computation.
+
+   Operationally `stop_gradient` is the identity function, that is, it returns
+   argument `x` unchanged. However, `stop_gradient` prevents the flow of
+   gradients during forward or reverse-mode automatic differentiation. If there
+   are multiple nested gradient computations, `stop_gradient` stops gradients
+   for all of them.
+
+   For example:
+
+   >>> jax.grad(lambda x: x**2)(3.)
+   array(6., dtype=float32)
+   >>> jax.grad(lambda x: jax.lax.stop_gradient(x)**2)(3.)
+   array(0., dtype=float32)
+   >>> jax.grad(jax.grad(lambda x: x**2))(3.)
+   array(2., dtype=float32)
+   >>> jax.grad(jax.grad(lambda x: jax.lax.stop_gradient(x)**2))(3.)
+   array(0., dtype=float32)
+   """
   return stop_gradient_p.bind(x)
 
 


Test patch:
Repository: google/jax

Issue:
Negative padding in pad_transpose
This is a request to support negative padding in the gradient of `lax.pad`, see: https://github.com/google/jax/blob/5592a063375a8feb4439e35edd8c335264dc1e6c/jax/lax.py#L2360


Discussion:

Resolving patch:
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -2356,20 +2356,16 @@ def _pad_shape_rule(operand, padding_value, padding_config):
 
 def _pad_transpose(t, operand, padding_value, padding_config):
   lo, hi, interior = zip(*padding_config)
-  if onp.any(onp.less(lo, 0)) or onp.any(onp.less(hi, 0)):
-    msg = "pad transpose not implemented for negative padding, got {}."
-    raise NotImplementedError(msg.format(padding_config))
 
   total = lambda x: _reduce_sum(x, list(range(t.ndim)))
 
-  t_op = lambda: slice(t, lo, onp.subtract(t.shape, hi), onp.add(interior, 1))
-  t_operand = t_op() if operand is None else None
+  def t_op():
+    unpad_config = zip(onp.negative(lo), onp.negative(hi), onp.zeros_like(interior))
+    unpadded = pad(t, onp.array(0., t.dtype), unpad_config)
+    return slice(unpadded, onp.zeros_like(lo), unpadded.shape, onp.add(interior, 1))
 
-  if padding_value is None:
-    t_operand = t_op() if t_operand is None else t_operand
-    t_padv = sub(total(t), total(t_operand))
-  else:
-    t_padv = None
+  t_operand = t_op() if operand is None else None
+  t_padv = sub(total(t), total(t_operand)) if padding_value is None else None
 
   return [t_operand, t_padv]
 


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -2007,7 +2007,7 @@ def testReshapeGrad(self, arg_shape, out_shape, dtype, rng):
        "shape": shape, "dtype": dtype, "pads": pads, "rng": jtu.rand_small()}
       for shape in [(2, 3)]
       for dtype in float_dtypes
-      for pads in [[(1, 2, 1), (0, 1, 0)]]))
+      for pads in [[(1, 2, 1), (0, 1, 0)], [(-1, 0, 0), (-1, 0, 2)]]))
   def testPadGrad(self, shape, dtype, pads, rng):
     tol = 1e-2 if onp.finfo(dtype).bits == 32 else None
Repository: google/jax

Issue:
[bug] using `lax.cond` with partial effect function gives TypeError
In the master branch, the following script shows the error
```sh
TypeError: could not merge true_fun and false_fun output pvals: ((ShapedArray(int32[]), None), JaxTuple(*,0)) and ((ShapedArray(int32[]),ShapedArray(int32[])), JaxTuple(*,*)).
```
```
lax.cond(True, (0, 0), lambda x: (x[0], 0), (1, 1), lambda x: x)

```

While these versions works well:
```
lax.cond(True, (0, 0), lambda x: (0, 0), (1, 1), lambda x: x)
```
or
```
lax.cond(True, (0, 0), lambda x: (x[0], x[1]), (1, 1), lambda x: x)
```


Discussion:

Resolving patch:
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -296,7 +296,7 @@ def join_pvals(pval1, pval2):
     pvals1, pvals2 = zip(pv1, const1), zip(pv2, const2)
     join_pvs, join_consts = unzip2(map(join_pvals, pvals1, pvals2))
     if all(isinstance(pv, AbstractValue) for pv in join_pvs):
-      return PartialVal(AbstractTuple(join_pvs), tuple(join_consts))
+      return PartialVal((AbstractTuple(join_pvs), tuple(join_consts)))
     else:
       return PartialVal((JaxprTracerTuple(join_pvs), tuple(join_consts)))
 


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1421,6 +1421,12 @@ def cfun(x):
     self.assertEqual(fun(4), cfun(4))
     self.assertEqual(cfun(4), (4, 2., 4.))
 
+  def testIssue514(self):
+    # just check this doesn't crash
+    lax.cond(True,
+            (0, 0), lambda x: (x[0], 0),
+            (1, 1), lambda x: x)
+
   def testScanAdd(self):
     def f(x, y):
       return x + y
Repository: google/jax

Issue:
undefined f if vmap applied to lu.WrappedFun
In `vmap`, `f` is not be defined whenever the if condition is false:

https://github.com/google/jax/blob/f8eda9bb645a1ed489f08be57bae380bbcf5fa0c/jax/api.py#L414-L420


Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -413,8 +413,7 @@ def vmap(fun, in_axes=0, out_axes=0):
 
   @wraps(fun, docstr=docstr)
   def batched_fun(*args, **kwargs):
-    if not isinstance(fun, lu.WrappedFun):
-      f = lu.wrap_init(fun, kwargs)
+    f = lu.wrap_init(fun, kwargs) if not isinstance(fun, lu.WrappedFun) else fun
     in_axes_ = in_axes if isinstance(in_axes, (list, tuple)) else (in_axes,) * len(args)
     in_flat, in_trees = unzip2(map(pytree_to_jaxtupletree, args))
     jaxtree_fun, out_tree = pytree_fun_to_jaxtupletree_fun(f, in_trees)


Test patch:
Repository: google/jax

Issue:
jax.config doesn't report command-line flags correctly
Another one from @jmgilmer and I - I don't think the jax config is parsing command line flags correctly.  I don't know if this is functionally important or just a reporting error - but it is certainly important for user scripts knowing what flags have been set.

If I run this script:
```
from absl import app, flags
from jax.config import config
FLAGS = flags.FLAGS
def main(_):
  print("FLAGS.jax_enable_x64", FLAGS.jax_enable_x64)
  print("FLAGS.jax_debug_nans", FLAGS.jax_debug_nans)
  print(config.values)
if __name__ == "__main__":
  config.config_with_absl()
  app.run(main)
```
I get the following problem: jax.config doesn't report the correct flag settings.
```
> python jaxtest.py --jax_enable_x64=1 --jax_debug_nans=1
FLAGS.jax_enable_x64 True
FLAGS.jax_debug_nans True
{'jax_enable_x64': 0, 'jax_xla_backend': 'xla', 'jax_backend_target': 'local', 'jax_platform_name': '', 'jax_device_values': 1, 'jax_debug_nans': 0, 'jax_disable_jit': 0}
```
if I run the same with envvars instead it works:
```
JAX_ENABLE_X64=1 JAX_DEBUG_NANS=1 python jaxtest.py
FLAGS.jax_enable_x64 True
FLAGS.jax_debug_nans True
{'jax_enable_x64': 1, 'jax_xla_backend': 'xla', 'jax_backend_target': 'local', 'jax_platform_name': '', 'jax_device_values': 1, 'jax_debug_nans': 1, 'jax_disable_jit': 0}
```

I've tried parsing the flags in different ways but nothing seems to fix the issue.


Discussion:

Resolving patch:
diff --git a/jax/config.py b/jax/config.py
--- a/jax/config.py
+++ b/jax/config.py
@@ -73,6 +73,8 @@ def config_with_absl(self):
       flag_type, meta_args, meta_kwargs = self.meta[name]
       absl_defs[flag_type](name, val, *meta_args, **meta_kwargs)
 
+    app.call_after_init(lambda: self.complete_absl_config(absl_flags))
+
   def complete_absl_config(self, absl_flags):
     for name, _ in self.values.items():
       self.update(name, getattr(absl_flags.FLAGS, name))
@@ -83,6 +85,7 @@ def parse_flags_with_absl(self):
       import absl.flags
       self.config_with_absl()
       absl.flags.FLAGS(sys.argv)
+      self.complete_absl_config(absl.flags)
       already_configured_with_absl = True
 
 


Test patch:
Repository: google/jax

Issue:
add jax.numpy.einsum_path
Requested by @murphyk.


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
@@ -1509,6 +1509,11 @@ def einsum(*operands):
   contractions = tuple(data[:3] for data in contractions)
   return _einsum(operands, contractions)
 
+@_wraps(onp.einsum_path)
+def einsum_path(subscripts, *operands, **kwargs):
+  optimize = kwargs.pop('optimize', 'greedy')
+  # using einsum_call=True here is an internal api for opt_einsum
+  return opt_einsum.contract_path(subscripts, *operands, optimize=optimize)
 
 @partial(jit, static_argnums=(1,))
 def _einsum(operands, contractions):


Test patch:
diff --git a/tests/lax_numpy_einsum_test.py b/tests/lax_numpy_einsum_test.py
--- a/tests/lax_numpy_einsum_test.py
+++ b/tests/lax_numpy_einsum_test.py
@@ -277,6 +277,22 @@ def test_ordered_front_batch_dim_case(self):
     s = 'ijkl,ijml->ijkm'
     self._check(s, x, y)
 
+  def test_einsum_path(self):
+    # just check examples from onp.einsum_path docstring
+    a = onp.random.rand(2, 2)
+    b = onp.random.rand(2, 5)
+    c = onp.random.rand(5, 2)
+
+    path_info = onp.einsum_path('ij,jk,kl->il', a, b, c, optimize='greedy')
+    self.assertEqual(str(path_info[0]), "['einsum_path', (1, 2), (0, 1)]")
+    self.assertEqual(path_info[1].split('\n')[0],
+                     '  Complete contraction:  ij,jk,kl->il')
+
+    # check this doesn't crash
+    I = onp.random.rand(10, 10, 10, 10)
+    C = onp.random.rand(10, 10)
+    onp.einsum_path('ea,fb,abcd,gc,hd->efgh', C, C, I, C, C, optimize='greedy')
+
 
 if __name__ == '__main__':
   absltest.main()
Repository: google/jax

Issue:
generators should not be accepted as values for the shape argument of numpy functions
```python3
import numpy as np
shape = (1 for _ in range(4)) 
np.zeros(shape)
```

results in `TypeError: expected sequence object with len >= 0 or a single integer`

whereas

```python3
import jax.numpy as np
shape = (1 for _ in range(4)) 
np.zeros(shape)
```

just runs through and produces `array([[[[0.]]]], dtype=float32)`.
This however can lead to bugs that are hard to find, because calling `np.zeros(shape)` a second time will create `array(0., dtype=float32)`

Conclusion: `jax.numpy` should behave like `numpy` and should also through a TypeError
generators should not be accepted as values for the shape argument of numpy functions
```python3
import numpy as np
shape = (1 for _ in range(4)) 
np.zeros(shape)
```

results in `TypeError: expected sequence object with len >= 0 or a single integer`

whereas

```python3
import jax.numpy as np
shape = (1 for _ in range(4)) 
np.zeros(shape)
```

just runs through and produces `array([[[[0.]]]], dtype=float32)`.
This however can lead to bugs that are hard to find, because calling `np.zeros(shape)` a second time will create `array(0., dtype=float32)`

Conclusion: `jax.numpy` should behave like `numpy` and should also through a TypeError


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
@@ -21,6 +21,7 @@
 import re
 import string
 import warnings
+import types
 
 import numpy as onp
 import opt_einsum
@@ -1221,6 +1222,8 @@ def full_like(a, fill_value, dtype=None):
 
 @_wraps(onp.zeros)
 def zeros(shape, dtype=onp.dtype("float64")):
+  if isinstance(shape, types.GeneratorType):
+    raise TypeError("expected sequence object with len >= 0 or a single integer")
   shape = (shape,) if onp.isscalar(shape) else shape
   return lax.full(shape, 0, dtype)
 


Test patch:
Repository: google/jax

Issue:
Regression in performance of samplers
With jax versions > 0.1.24, random samplers are slow. I think that it is due to recent changes on how jit works with static_args/kwargs. A script to reproduce,
```
import time
from jax import random

t = time.time()
random.normal(random.PRNGKey(0), shape=(1,))
print(time.time() - t)
t = time.time()
random.normal(random.PRNGKey(0), shape=(1,))
print(time.time() - t)
```
which returns `0.12923526763916016` and `0.12221717834472656`.

However, if we wrap these samplers in some function, then it is fast. For example,
```
def f():
    return random.normal(random.PRNGKey(0), shape=(1,))

t = time.time()
f()
print(time.time() - t)
t = time.time()
f()
print(time.time() - t)
```
will return `0.12787413597106934` and `0.0010831356048583984`.

I think that there is a small bug elsewhere which forces the sampler recompile. If this is an expected behaviour, then which function we should use to wrap these samplers to make it not recompiled?
cc @neerajprad



Discussion:
Thanks so much for catching this! We will fix it asap.

(Longer term we aim to set up a continuous benchmarking solution to avoid regressions.)

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -46,7 +46,7 @@
                         tree_map, tree_flatten, tree_unflatten, tree_structure,
                         tree_transpose, leaf)
 from .util import (unzip2, unzip3, curry, partial, safe_map, safe_zip,
-                   WrapHashably, prod)
+                   WrapHashably, Hashable, prod)
 from .lib.xla_bridge import canonicalize_dtype, device_count
 from .abstract_arrays import ShapedArray
 from .interpreters import partial_eval as pe
@@ -824,11 +824,19 @@ def _argnums_partial(f, dyn_argnums, args):
     dyn_argnums = (dyn_argnums,)
   else:
     dyn_argnums = tuple(dyn_argnums)
-  fixed_args = tuple([None if i in dyn_argnums else WrapHashably(arg)
+  fixed_args = tuple([None if i in dyn_argnums else _wrap_hashably(arg)
                       for i, arg in enumerate(args)])
   dyn_args = tuple(args[i] for i in dyn_argnums)
   return _argnums_partial_(f, dyn_argnums, fixed_args), dyn_args
 
+def _wrap_hashably(arg):
+  try:
+    hash(arg)
+  except TypeError:
+    return WrapHashably(arg)
+  else:
+    return Hashable(arg)
+
 @lu.transformation
 def _argnums_partial_(dyn_argnums, fixed_args, *dyn_args, **kwargs):
   args = [None if arg is None else arg.val for arg in fixed_args]
diff --git a/jax/util.py b/jax/util.py
--- a/jax/util.py
+++ b/jax/util.py
@@ -189,6 +189,8 @@ def prod(xs):
 
 
 class WrapHashably(object):
+  __slots__ = ["val"]
+
   def __init__(self, val):
     self.val = val
 
@@ -198,6 +200,18 @@ def __hash__(self):
   def __eq__(self, other):
     return self.val is other.val
 
+class Hashable(object):
+  __slots__ = ["val"]
+
+  def __init__(self, val):
+    self.val = val
+
+  def __hash__(self):
+    return hash(self.val)
+
+  def __eq__(self, other):
+    return self.val == other.val
+
 
 
 def get_module_functions(module):


Test patch:
Repository: google/jax

Issue:
Scalars passed into np.array should return 0-dim arrays
import jax.numpy as np
import numpy as onp

onp.array(0).dtype # works, because onp.array(0) returns a 0-dim array
np.array(0).dtype # doesn't work, because np.array(0) returns an int

This showed up when trying to do `lax.dynamic_update_slice(array, value, np.array(0))`


Discussion:
#124  says this was fixed, but I still get an error:

```
In [1]: import numpy as onp

In [2]: import jax.numpy as np

In [3]: onp.array(0).dtype
Out[3]: dtype('int64')

In [4]: np.array(0).dtype
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-4-05f7ef02f43f> in <module>
----> 1 np.array(0).dtype

AttributeError: 'int' object has no attribute 'dtype'
```

I'm on commit 6476d5ffc6bee0e8d9d7b54914565d798b8dda29.
I am not sure if this is a regression, but we noticed this recently as well when we found that `scan` does not work well with scalar arrays (https://github.com/pyro-ppl/numpyro/pull/91). The reason is that `jax.numpy.array(1.)` returns a python float instead of a `DeviceArray` unlike numpy. 

cc. @mattjj in case this needs to be reopened.
@kroq-gar78 that PR was closed but not merged, so this bug still isn't fixed and remains open.

@neerajprad we're going to check in a from-scratch rewrite of `scan` in the next week or two, though that new version will need testing and could have similar issues at first. Thanks for pointing that out!

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -514,7 +514,7 @@ def reshape(operand, new_sizes, dimensions=None):
   """
   same_shape = onp.shape(operand) == tuple(new_sizes)
   same_dims = dimensions is None or tuple(dimensions) == tuple(range(onp.ndim(operand)))
-  if same_shape and same_dims:
+  if onp.shape(operand) and same_shape and same_dims:
     return operand
   else:
     return reshape_p.bind(


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
@@ -1051,7 +1051,7 @@ def testAverage(self, shape, dtype, axis, weights_shape, returned, rng):
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_arg{}".format(i), "arg": arg}
       for i, arg in enumerate([
-          [1, 2, 3], [1., 2., 3.],
+          3., [1, 2, 3], [1., 2., 3.],
           [[1, 2], [3, 4], [5, 6]], [[1, 2.], [3, 4], [5, 6]],
           [[3, onp.array(2), 1], onp.arange(3.)],
       ])))
@@ -1060,6 +1060,9 @@ def testArray(self, arg):
     self._CheckAgainstNumpy(onp.array, lnp.array, args_maker, check_dtypes=True)
     self._CompileAndCheck(lnp.array, args_maker, check_dtypes=True)
 
+  def testIssue121(self):
+    assert not onp.isscalar(lnp.array(3))
+
   def testArrayMethod(self):
     class arraylike(object):
       dtype = onp.float32