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

Issue:
pmap error message for mismatched rank with PyTree args
Small usability request :)

Currently, if you pass in an argument with an incorrect rank to a pmap you get the error: `ValueError: pmap got arg {arg_num} of rank 0 but axis to be mapped 0`

The `arg_num` index will correctly point to the position of the faulty argument, but only if your arguments do not contain PyTrees. If your arguments contain PyTrees you could for example get a `ValueError: pmap got arg 35 of rank 0 but axis to be mapped 0` on a function with 5 arguments.

If you pass in an incorrectly sized argument, there is a separate error message for PyTree args (https://github.com/google/jax/blob/ee9ca569890fae79c7e82fe5bad26e5ce9e72227/jax/api.py#L1275).   Maybe something similar can be done here for an incorrect rank?


Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -1234,14 +1234,15 @@ def batched_fun(*args):
 
   return batched_fun
 
-def _get_axis_size(name: str, i:int, shape: Tuple[int, ...], axis: int):
-  try:
-    return shape[axis]
-  except (IndexError, TypeError) as e:
-    raise ValueError(f"{name} got arg {i} of rank {len(shape)} "
-                     f"but axis to be mapped {axis}") from e
-
 def _mapped_axis_size(tree, vals, dims, name):
+  def _get_axis_size(name: str, i:int, shape: Tuple[int, ...], axis: int):
+    try:
+      return shape[axis]
+    except (IndexError, TypeError) as e:
+      ranks = tree_unflatten(tree, [np.ndim(x) for x, d in zip(vals, dims)])
+      raise ValueError(f"{name} got arg {i} of rank {len(shape)} but axis to be mapped {axis}. "
+                       f"The tree of ranks is:\n{ranks}") from e
+
   mapped_axis_sizes = {_get_axis_size(name, i, np.shape(x), d)
                        for i, (x, d) in enumerate(zip(vals, dims))
                        if d is not None}


Test patch:
Repository: google/jax

Issue:
docs: jax.scipy.stats.poisson not on the public API docs webpage
Hi. With the release of JAX [`v0.2.4`](https://github.com/google/jax/releases/tag/jax-v0.2.4) (congrats on the release btw) we noticed that `jax.scipy.stats.poisson` is not documented on the [public API docs for the `jax.scipy` package](https://jax.readthedocs.io/en/stable/jax.scipy.html#jax-scipy-stats). However, other distributions like [`jax.scipy.stats.norm` are documented](https://jax.readthedocs.io/en/stable/jax.scipy.html#module-jax.scipy.stats.norm).

Is this intentional (are there upcoming plans for `jax.scipy.stats.poisson`)? Or does `jax.scipy.stats.poisson` just still need to get put into Sphinx?

(cc @lukasheinrich @kratsg)


Discussion:

Resolving patch:
diff --git a/jax/scipy/stats/__init__.py b/jax/scipy/stats/__init__.py
--- a/jax/scipy/stats/__init__.py
+++ b/jax/scipy/stats/__init__.py
@@ -14,17 +14,17 @@
 
 # flake8: noqa: F401
 from . import bernoulli
-from . import poisson
 from . import beta
 from . import cauchy
 from . import dirichlet
 from . import expon
 from . import gamma
+from . import geom
 from . import laplace
+from . import logistic
 from . import multivariate_normal
 from . import norm
 from . import pareto
+from . import poisson
 from . import t
 from . import uniform
-from . import logistic
-from . import geom


Test patch:
Repository: google/jax

Issue:
jax.scipy.minimize optional function args not being passed properly
In `jax.scipy.optimize._minimize`,
```python
fun_with_args = partial(fun, *args)
```
incorrectly puts the `args` before the parameter when `args` are provided.
E.g. 

```python
from functools import partial
func = lambda a,b: print(a,b)

args = ('b',)
partial(func, *args)('a')
# b a instead of a b as intended
```
jax.scipy.minimize optional function args not being passed properly
In `jax.scipy.optimize._minimize`,
```python
fun_with_args = partial(fun, *args)
```
incorrectly puts the `args` before the parameter when `args` are provided.
E.g. 

```python
from functools import partial
func = lambda a,b: print(a,b)

args = ('b',)
partial(func, *args)('a')
# b a instead of a b as intended
```


Discussion:

Resolving patch:
diff --git a/jax/_src/scipy/optimize/line_search.py b/jax/_src/scipy/optimize/line_search.py
--- a/jax/_src/scipy/optimize/line_search.py
+++ b/jax/_src/scipy/optimize/line_search.py
@@ -112,7 +112,7 @@ def body(state):
 
     # This will cause the line search to stop, and since the Wolfe conditions
     # are not satisfied the minimization should stop too.
-    state = state._replace(failed=state.failed | (dalpha <= 1e-5))
+    state = state._replace(failed=state.failed | (dalpha <= 1e-10))
 
     # Cubmin is sometimes nan, though in this case the bounds check will fail.
     a_j_cubic = _cubicmin(state.a_lo, state.phi_lo, state.dphi_lo, state.a_hi,
diff --git a/jax/_src/scipy/optimize/minimize.py b/jax/_src/scipy/optimize/minimize.py
--- a/jax/_src/scipy/optimize/minimize.py
+++ b/jax/_src/scipy/optimize/minimize.py
@@ -11,9 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from functools import partial
 from typing import Any, Callable, Mapping, Optional, Tuple, Union
-
 from .bfgs import minimize_bfgs
 from typing import NamedTuple
 import jax.numpy as jnp
@@ -92,7 +90,7 @@ def minimize(
   if options is None:
     options = {}
 
-  fun_with_args = partial(fun, *args)
+  fun_with_args = lambda x: fun(x, *args)
 
   if method.lower() == 'bfgs':
     results = minimize_bfgs(fun_with_args, x0, **options)


Test patch:
diff --git a/tests/scipy_optimize_test.py b/tests/scipy_optimize_test.py
--- a/tests/scipy_optimize_test.py
+++ b/tests/scipy_optimize_test.py
@@ -86,6 +86,14 @@ def min_op(x0):
     scipy_res = scipy.optimize.minimize(func(np), x0, method='BFGS').x
     self.assertAllClose(scipy_res, jax_res, atol=2e-5, check_dtypes=False)
 
+  def test_fixes4594(self):
+    n = 2
+    A = jnp.eye(n) * 1e4
+    def f(x):
+      return jnp.mean((A @ x) ** 2)
+    results = jax.scipy.optimize.minimize(f, jnp.ones(n), method='BFGS')
+    self.assertAllClose(results.x, jnp.zeros(n), atol=1e-6, rtol=1e-6)
+
 
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax

Issue:
_min = builtins.max ?
There is a very suspicious assignment `_min = builtins.max` in `lax.py`. Is this a potential bug or is there actually a reason for this?

https://github.com/google/jax/blob/5c9204a4613b5a92dc6c60c13579d783bab4976f/jax/_src/lax/lax.py#L58


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -55,7 +55,7 @@
 FLAGS = flags.FLAGS
 
 _max = builtins.max
-_min = builtins.max
+_min = builtins.min
 _reduce = functools.reduce
 
 Array = Any


Test patch:
Repository: google/jax

Issue:
Adam turns Float32 into Float64
Dear Jax developers, 

Thanks for the awesome work! When I set `jax_enable_x64=True`, I found the Adam optimizer turns float32 into float64 after one update. This happens for adam exclusively, not for sgd or rmsprop. Is this a known issue or a specific design? Do we have a nice way to counter this ? 

I am using `jax==0.2.1` and `jaxlib==0.1.55`. I run the following code in purely cpu, though commenting `jit` will remove this phenomenon.  The results are 
```
0 float32
1 float64
2 float64
```

```

import numpy as onp

import jax.numpy as jnp
from jax.config import config
config.update('jax_enable_x64', True)
from jax import grad, jit
from jax.experimental import optimizers


x = onp.random.RandomState(0).normal(size=[10, 10]).astype('float32')

opt_init, opt_update, get_params = optimizers.adam(1e-3)
# opt_init, opt_update, get_params = optimizers.sgd(1e-3)
# opt_init, opt_update, get_params = optimizers.rmsprop(1e-3)
opt_state = opt_init(jnp.asarray(x).astype('float32'))

@jit
def loss(x):
    return jnp.mean(x ** 2.)

@jit
def update(i, opt_state):
    x = get_params(opt_state)
    return opt_update(i, grad(loss)(x), opt_state)

print(0, x.dtype)

opt_state = update(0, opt_state)
x = get_params(opt_state)
print(1, x.dtype)

opt_state = update(1, opt_state)
x = get_params(opt_state)
print(2, x.dtype)

```



Discussion:

Resolving patch:
diff --git a/jax/experimental/optimizers.py b/jax/experimental/optimizers.py
--- a/jax/experimental/optimizers.py
+++ b/jax/experimental/optimizers.py
@@ -406,8 +406,8 @@ def update(i, g, state):
     x, m, v = state
     m = (1 - b1) * g + b1 * m  # First  moment estimate.
     v = (1 - b2) * jnp.square(g) + b2 * v  # Second moment estimate.
-    mhat = m / (1 - b1 ** (i + 1))  # Bias correction.
-    vhat = v / (1 - b2 ** (i + 1))
+    mhat = m / (1 - jnp.asarray(b1, m.dtype) ** (i + 1))  # Bias correction.
+    vhat = v / (1 - jnp.asarray(b2, m.dtype) ** (i + 1))
     x = x - step_size(i) * mhat / (jnp.sqrt(vhat) + eps)
     return x, m, v
   def get_params(state):
@@ -442,7 +442,8 @@ def update(i, g, state):
     x, m, u = state
     m = (1 - b1) * g + b1 * m  # First  moment estimate.
     u = jnp.maximum(b2 * u, jnp.abs(g))  # Update exponentially weighted infinity norm.
-    x = x - (step_size(i) / (1 - b1 ** (i + 1))) * m / (u + eps)
+    x = (x - (step_size(i) / (1 - jnp.asarray(b1, m.dtype) ** (i + 1))) * m
+         / (u + eps))
     return x, m, u
   def get_params(state):
     x, _, _ = state


Test patch:
Repository: google/jax

Issue:
Add Python 3.9 wheels for jaxlib
Sorry for the bother, but I'm having some problems with installing jaxlib in Python 3.9:
```
> pip install jaxlib==0.1.55
ERROR: Could not find a version that satisfies the requirement jaxlib==0.1.55 (from versions: 0.1, 0.1.1, 0.1.4)
ERROR: No matching distribution found for jaxlib==0.1.55
```
@hawkinsp 


Discussion:
We haven't released wheels for Python 3.9 yet. But we can try to do that! One slight blocker is that NumPy and SciPy haven't released 3.9 wheels yet (although we can most likely build them from source). So I'd be tempted to wait a week or two before trying to release 3.9 wheels.
@hawkinsp No rush!  Just thought it might have been an oversight since a few versions are mysteriously there.
@hawkinsp we would like to build jaxlib from source in conda-forge, at the moment we are re-packaging the wheel. If there is any interested from the developers in doing so please open an issue in https://github.com/conda-forge/jaxlib-feedstock, we'd love to have some help!

PS: we already have scipy and numpy on py39!
FYI:  [NumPy](https://pypi.org/project/numpy/#files) and [SciPy](https://pypi.org/project/scipy/#files) have now released wheels for Python 3.9.

Resolving patch:
diff --git a/build/build.py b/build/build.py
--- a/build/build.py
+++ b/build/build.py
@@ -179,6 +179,7 @@ def check_bazel_version(bazel_path, min_version, max_version):
 common --experimental_repo_remote_exec
 
 build --repo_env PYTHON_BIN_PATH="{python_bin_path}"
+build --action_env=PYENV_ROOT
 build --python_path="{python_bin_path}"
 build --repo_env TF_NEED_CUDA="{tf_need_cuda}"
 build --action_env TF_CUDA_COMPUTE_CAPABILITIES="{cuda_compute_capabilities}"


Test patch:
Repository: google/jax

Issue:
Error caused by shutil.rmtree 
```
Traceback (most recent call last):
  File "\\?\C:\Users\cloud\AppData\Local\Temp\Bazel.runfiles_vfpgffuf\runfiles\__main__\build\install_xla_in_source_tree.py", line 83, in <module>
    shutil.rmtree(jaxlib_dir)
  File "C:\Users\cloud\miniconda3\lib\shutil.py", line 516, in rmtree
    return _rmtree_unsafe(path, onerror)
  File "C:\Users\cloud\miniconda3\lib\shutil.py", line 400, in _rmtree_unsafe
    onerror(os.unlink, fullname, sys.exc_info())
  File "C:\Users\cloud\miniconda3\lib\shutil.py", line 398, in _rmtree_unsafe
    os.unlink(fullname)
WindowsError: [Error 5] Access is denied.: 'D:\\jax\\build\\jaxlib\\cublas_kernels.pyd'
```

This only happens on rebuild.

The reason is `shutil.rmtree` will not delete readonly file on Windows.


Discussion:
@cloudhan can you try with https://github.com/google/jax/pull/4982 patched in?

This reworks the build a bit more so it builds a wheel file rather than writing into the source tree. Since it uses a temporary directory it doesn't need this particular line of code.
I just updated that PR to fix some more windows-specific problems and verified I was able to build a windows wheel.

Resolving patch:
diff --git a/build/build_wheel.py b/build/build_wheel.py
--- a/build/build_wheel.py
+++ b/build/build_wheel.py
@@ -57,13 +57,19 @@ def _copy_so(src_file, dst_dir, dst_filename=None):
     else:
       dst_filename = src_filename
   dst_file = os.path.join(dst_dir, dst_filename)
-  shutil.copy(src_file, dst_file)
+  if _is_windows():
+    shutil.copyfile(src_file, dst_file)
+  else:
+    shutil.copy(src_file, dst_file)
 
 
 def _copy_normal(src_file, dst_dir, dst_filename=None):
   src_filename = os.path.basename(src_file)
   dst_file = os.path.join(dst_dir, dst_filename or src_filename)
-  shutil.copy(src_file, dst_file)
+  if _is_windows():
+    shutil.copyfile(src_file, dst_file)
+  else:
+    shutil.copy(src_file, dst_file)
 
 
 def copy_file(src_file, dst_dir, dst_filename=None):
@@ -169,4 +175,3 @@ def build_wheel(sources_path, output_path):
 finally:
   if tmpdir:
     tmpdir.cleanup()
-


Test patch:
Repository: google/jax

Issue:
repr(sharded_device_array) should print ShardedDeviceArray
Just noticed that we had a trivial regression in how repr(sharded_device_array) works:

```python
In [1]: from jax.api import device_put_sharded, devices

In [2]: device_put_sharded([1], devices())
Out[2]: DeviceArray([1], dtype=int32)
```

We use repr(x) to show the type of x, e.g. in demos, so we probably want to put this back.


Discussion:
Ah, [it's just this change](https://github.com/google/jax/commit/ad2de7554641d5ee986f43837bacaebc79c2337b#diff-121975a1a88d61a259f8423655e23a1b389eb18c399b267395d2aba5a196bae9R1170). @jblespiau to answer your question in that TODO, yes we should put it back ;)

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -1167,8 +1167,7 @@ def copy(self):
 
   def __repr__(self):
     line_width = np.get_printoptions()["linewidth"]
-    # TODO(jblespia): Should we put back self.__class__.__name__ ?
-    prefix = '{}('.format("DeviceArray")
+    prefix = '{}('.format(self.__class__.__name__.lstrip('_'))
     s = np.array2string(self._value, prefix=prefix, suffix=',',
                         separator=', ', max_line_width=line_width)
     dtype_str = 'dtype={})'.format(self.dtype.name)


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -2091,6 +2091,10 @@ def test_device_put_replicated_pytree(self):
     self.assertTrue(all(b.device() == d for b, d in zip(y2.device_buffers, devices)))
     self.assertArraysEqual(y2, np.stack([xs['b'] for _ in devices]))
 
+  def test_repr(self):
+    x = jax.device_put_replicated(1, jax.devices())
+    self.assertStartsWith(repr(x), 'ShardedDeviceArray')
+
 
 class SpecToIndicesTest(jtu.JaxTestCase):
Repository: google/jax

Issue:
jnp.acosh and np.acosh follow different conventions
The following code outputs different results

```python
import numpy as np
from jax import numpy as jnp

x = -4.69577+2.1128223j

print(np.arccosh(x))
print(jnp.arccosh(x))
```

While there are infinitely many possible results, numpy defines a convention in its documentation for what the result should be:
https://numpy.org/doc/stable/reference/generated/numpy.arccosh.html

I do not know if it is that important, but figured I would report since we expect `jnp` and `np` to have the same behavior :)


Discussion:
Thanks! We should change JAX to match numpy's convention here.

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
@@ -417,7 +417,6 @@ def fn(x1, x2):
 arcsinh = _one_to_one_unop(np.arcsinh, lax.asinh, True)
 tanh = _one_to_one_unop(np.tanh, lax.tanh, True)
 arcsinh = _one_to_one_unop(np.arcsinh, lax.asinh, True)
-arccosh = _one_to_one_unop(np.arccosh, lax.acosh, True)
 arctanh = _one_to_one_unop(np.arctanh, lax.atanh, True)
 sqrt = _one_to_one_unop(np.sqrt, lax.sqrt, True)
 
@@ -437,6 +436,14 @@ def fn(x1, x2):
 float_power = _one_to_one_binop(np.float_power, lax.pow, True)
 nextafter = _one_to_one_binop(np.nextafter, lax.nextafter, True, True)
 
+@_wraps(np.arccosh)
+def arccosh(x):
+  # Note: arccosh is multi-valued for complex input, and lax.acosh uses a different
+  # convention than np.arccosh.
+  out = lax.acosh(*_promote_args_inexact("arccosh", x))
+  if issubdtype(out.dtype, np.complexfloating):
+    out = where(real(out) < 0, lax.neg(out), out)
+  return out
 
 def _comparison_op(numpy_fn, lax_fn):
   def fn(x1, x2):


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
@@ -180,10 +180,10 @@ def op_record(name, nargs, dtypes, shapes, rng_factory, diff_modes,
               inexact=True),
     op_record("arctan2", 2, float_dtypes, all_shapes, jtu.rand_small, ["rev"],
               inexact=True),
-    op_record("arcsinh", 1, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],
-              inexact=True),
-    op_record("arccosh", 1, number_dtypes, all_shapes, jtu.rand_positive, ["rev"],
-              inexact=True),
+    op_record("arcsinh", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
+              inexact=True, tolerance={np.complex64: 2E-4, np.complex128: 2E-14}),
+    op_record("arccosh", 1, number_dtypes, all_shapes, jtu.rand_default, ["rev"],
+              inexact=True, tolerance={np.complex64: 2E-2, np.complex128: 2E-12}),
     op_record("arctanh", 1, number_dtypes, all_shapes, jtu.rand_small, ["rev"],
               inexact=True, tolerance={np.float64: 1e-9}),
 ]
Repository: google/jax

Issue:
Composition of jet with grad
@mattjj 
Hi all,

I noticed that sometimes composing jet with grad produces an error. See example below:

```
import jax.numpy as jnp
from jax import grad
from jax.experimental.jet import jet

f = lambda x, ε: x * ε + ε + x

def g(ε):
  x = jnp.array(1.)
  return grad(f)(x, ε)

jet(g, (1.,), ([1.],))
```
Any insights if this is a bug?

Thanks,





Discussion:
Thanks for raising this! It's a missing rule for `jet` (which needs one rule for each primitive), so it's more of an 'enhancement needed' rather than a 'bug'. Or maybe I just prefer that positive framing of things :)

Resolving patch:
diff --git a/jax/experimental/jet.py b/jax/experimental/jet.py
--- a/jax/experimental/jet.py
+++ b/jax/experimental/jet.py
@@ -229,6 +229,7 @@ def linear_prop(prim, primals_in, series_in, **params):
 deflinear(lax.conj_p)
 deflinear(lax.imag_p)
 deflinear(lax.add_p)
+deflinear(ad_util.add_jaxvals_p)
 deflinear(lax.sub_p)
 deflinear(lax.convert_element_type_p)
 deflinear(lax.broadcast_in_dim_p)


Test patch:
diff --git a/tests/jet_test.py b/tests/jet_test.py
--- a/tests/jet_test.py
+++ b/tests/jet_test.py
@@ -19,6 +19,7 @@
 import numpy as np
 import unittest
 
+import jax
 from jax import test_util as jtu
 import jax.numpy as jnp
 import jax.scipy.special
@@ -379,6 +380,14 @@ def g(x):
     assert g_out_primals == f_out_primals
     assert g_out_series == f_out_series
 
+  def test_add_any(self):
+    # https://github.com/google/jax/issues/5217
+    f = lambda x, eps: x * eps + eps + x
+    def g(eps):
+      x = jnp.array(1.)
+      return jax.grad(f)(x, eps)
+    jet(g, (1.,), ([1.],))  # doesn't crash
+
 
 if __name__ == '__main__':
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax

Issue:
remove tie-in operands under omnistaging for primitives that don't need them
With omnistaging enabled, any primitive such as `create_token`, and its corresponding `lax` function, no longer needs to take a tie-in argument for data dependence.

Removing `create_token`'s operands also enforces that the primitive can't be linear, so it can't be transposed, and the system shouldn't try to transpose it. As @mattjj points out, a nullary primitive wouldn't appear in a jaxpr to be transposed because partial evaluation will never pick it up. This ought to address [this comment in #4292](https://github.com/google/jax/issues/4292#issuecomment-746289386).


Discussion:

Resolving patch:
diff --git a/jax/_src/lax/lax.py b/jax/_src/lax/lax.py
--- a/jax/_src/lax/lax.py
+++ b/jax/_src/lax/lax.py
@@ -5898,21 +5898,25 @@ def _stop_gradient_batch_rule(batched_args, batch_dims):
 batching.primitive_batchers[ad_util.stop_gradient_p] = _stop_gradient_batch_rule
 
 
-def create_token(x):
+def create_token(_=None):
   """Creates an XLA token value with no preconditions for sequencing effects.
 
   Experimental.
 
-  Args:
-    x: a dummy argument used to tie the CreateToken operator into a trace. The
-       value of `x` is ignored.
+  The argument is ignored. It exists for backward compatibility.
   """
-  # x is a dummy argument used to tie the operator into a trace.
-  return create_token_p.bind(stop_gradient(x))
+  if config.omnistaging_enabled:
+    return create_token_p.bind()
+  else:
+    x = _
+    if x is None:
+      raise ValueError(
+          'create_token needs a tie-in operand unless omnistaging is enabled.')
+    return create_token_p.bind(stop_gradient(x))
 
 create_token_p = Primitive("create_token")
 create_token_p.def_impl(partial(xla.apply_primitive, create_token_p))
-create_token_p.def_abstract_eval(lambda _: abstract_token)
+create_token_p.def_abstract_eval(lambda *_: abstract_token)
 xla.translations[create_token_p] = lambda c, *_: xops.CreateToken(c)
 
 def after_all(*operands):


Test patch:
Repository: google/jax

Issue:
jax.partial breaks signature inspection
```python
import inspect
from jax import partial

def f(x, y):
  return x + y

add_one = partial(f, y=1)
inspect.signature(add_one).bind(1)
# TypeError: missing a required argument: 'y'
```
Compare to `functools.partial`:
```python
import inspect
from functools import partial

def f(x, y):
  return x + y

add_one = partial(f, y=1)
inspect.signature(add_one).bind(1)
# <BoundArguments (x=1)>
```
Proposed fix is to alias `jax.partial` to `functools.partial`, and possibly one day deprecate it.


Discussion:

Resolving patch:
diff --git a/jax/util.py b/jax/util.py
--- a/jax/util.py
+++ b/jax/util.py
@@ -21,6 +21,7 @@
 
 import numpy as np
 
+partial = functools.partial
 
 def safe_zip(*args):
   n = len(args[0])
@@ -90,12 +91,6 @@ def split_dict(dct, names):
 def concatenate(xs):
   return list(it.chain.from_iterable(xs))
 
-def partial(fun, *args, **kwargs):
-  wrapped = functools.partial(fun, *args, **kwargs)
-  functools.update_wrapper(wrapped, fun)
-  wrapped._bound_args = args
-  return wrapped
-
 class partialmethod(functools.partial):
   def __get__(self, instance, owner):
     if instance is None:


Test patch:
Repository: google/jax

Issue:
PEP 561 Compliance
It's really great to see type annotations throughout the library! When I try to use `mypy` on my project which uses JAX, I get the following error:

```bash
 error: Skipping analyzing 'jax': found module but no type hints or library stubs
```
I think there needs to be an empty `py.typed` file in the root of the source files, as per [PEP 561](https://www.python.org/dev/peps/pep-0561/), and a corresponding entry in the `setup.py`: 
```python
setup(
  ...
  package_data={'jax': ['py.typed']}
  ...
}
```


Discussion:
This is in progress in #4711
Thanks @jakevdp!

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -26,6 +26,7 @@
     author='JAX team',
     author_email='jax-dev@google.com',
     packages=find_packages(exclude=["examples"]),
+    package_data={'jax': ['py.typed']},
     python_requires='>=3.6',
     install_requires=[
         'numpy >=1.12',


Test patch: