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

Issue:
Documentation on how to vmap on dictionary entries
I have a related question to #2367: Is it possible to use `vmap` on an array that is in a dictionary? And if yes, how?

To stick with the example given in #2367, how to `vmap` 'b', which is an array within a dictionary:
```
import jax.numpy as np
from jax import vmap

dictionary = {'a': 5., 'b': np.arange(5)}
c = 1.
d = 2.

def f(dct, c, d):
  return dct['a'] + dct['b'] + c + d

result = vmap(f, magic)(dictionary, c, d)
```

I don't understand how I would have to create the magic axes tuple for this example or if it is even possible? I think it would be great if there would be a related example in the [documentation of vmap](https://jax.readthedocs.io/en/latest/jax.html?highlight=vmap#jax.vmap).


Discussion:
Yes, set `magic = ({'a': None, 'b': 0}, None, None)`, like this:

```python
import jax.numpy as np
from jax import vmap

dictionary = {'a': 5., 'b': np.arange(5)}
c = 1.
d = 2.

def f(dct, c, d):
  return dct['a'] + dct['b'] + c + d

result = vmap(f, in_axes=({'a': None, 'b': 0}, None, None))(dictionary, c, d)
```

This sentence in the vmap docstring is intended to describe the behavior:

> If the positional arguments to fun are container types, the corresponding element of in_axes can itself be a matching container, so that distinct array axes can be mapped for different container elements. in_axes must be a container tree prefix of the positional argument tuple passed to fun.

Here, the argument in question is a dictionary, so we make the corresponding entry of `in_axes` a dictionary too (with the same keys).

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -796,6 +796,17 @@ def vmap(fun: Callable, in_axes=0, out_axes=0) -> Callable:
   >>> print(vfoo(tree).shape)
   (6, 2, 5)
 
+  Here's another example using container types in ``in_axes``, this time a
+  dictionary, to specify the elements of the container to map over:
+
+  >>> dct = {'a': 0., 'b': np.arange(5.)}
+  >>> x = 1.
+  >>> def foo(dct, x):
+  ...  return dct['a'] + dct['b'] + x
+  >>> out = vmap(foo, in_axes=({'a': None, 'b': 0}, None))(dct, x)
+  >>> print(out)
+  [1. 2. 3. 4. 5.]
+
   The results of a vectorized function can be mapped or unmapped.
   For example, the function below returns a pair with the first
   element mapped and the second unmapped. Only for unmapped results


Test patch:
Repository: google/jax

Issue:
Cast breaks jit
Casting via `jnp.float32(x)` breaks jit. 

> *** Exception: Tracer can’t be used with raw numpy functions. You might have
> import numpy as np
> instead of
> import jax.numpy as jnp

The workaround is to use `jnp.asarray(x, dtype=jnp.float32)` instead, but the error is very confusing and it would be nice if this was done automatically under the hood by jax.


Discussion:
I suspect to fix this all we need to do is change this line of code:
https://github.com/google/jax/blob/bb2127cebd3ade161109ee4919a92aaff5c788c1/jax/numpy/lax_numpy.py#L151
to do something different for JAX array types.

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
@@ -147,7 +147,7 @@ def __ne__(self, other):
     return not (self == other)
 
   def __call__(self, x):
-    return array(self.dtype.type(x), dtype=self.dtype)
+    return array(x, dtype=self.dtype)
 
 def _make_scalar_type(np_scalar_type):
   return _ScalarMeta(np_scalar_type.__name__, (object,),


Test patch:
diff --git a/tests/dtypes_test.py b/tests/dtypes_test.py
--- a/tests/dtypes_test.py
+++ b/tests/dtypes_test.py
@@ -173,5 +173,11 @@ class AnEnum(enum.IntEnum):
     np.testing.assert_equal(np.int32(101), np.int32(AnEnum.B))
     np.testing.assert_equal(jnp.int32(101), jnp.int32(AnEnum.B))
 
+  def testScalarCastInsideJitWorks(self):
+    # jnp.int32(tracer) should work.
+    self.assertEqual(jnp.int32(101),
+                     jax.jit(lambda x: jnp.int32(x))(jnp.float32(101.4)))
+
+
 if __name__ == "__main__":
   absltest.main()
Repository: google/jax

Issue:
Incorrect output from jax.numpy.column_stack
Using the CPU-only version of JAX (`jax 0.1.69`, `jaxlib 0.1.47`) in a fresh conda environment.

```
import numpy as np
import jax.numpy as jp

x = np.array([1., 1.])
y = np.array([2., 2.])

print(np.column_stack((x, y)))   # [[1. 2.], [1. 2.]], expected
print(jp.column_stack((x, y)))   # [[1., 1., 2., 2.]], incorrect
```


Discussion:
Hmm seems close enough up to numerical error...

Just kidding! Thanks for spotting the issue. Our tests must be lacking here.
Tests? Who needs tests? :grin: https://github.com/google/jax/search?q=column_stack

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
@@ -2030,7 +2030,7 @@ def column_stack(tup):
   for v in tup:
     arr = array(v)
     if arr.ndim < 2:
-      arr = expand_dims(arr, axis=0)
+      arr = atleast_2d(arr).T
     arrays.append(arr)
   return concatenate(arrays, 1)
 


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
@@ -1659,6 +1659,26 @@ def testDigitize(self, xshape, binshape, right, reverse, dtype, rng_factory):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": "_{}".format(
+          jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes)),
+       "shape": shape, "dtypes": dtypes}
+      for dtypes in [
+        [np.float32],
+        [np.float32, np.float32],
+        [np.float32, np.int32, np.float32],
+        [np.float32, np.int64, np.float32],
+        [np.float32, np.int32, np.float64],
+      ]
+      for shape in [(), (2,), (3, 4), (1, 5)]))
+  def testColumnStack(self, shape, dtypes):
+    rng = jtu.rand_default(self.rng())
+    args_maker = lambda: [[rng(shape, dtype) for dtype in dtypes]]
+    np_fun = _promote_like_jnp(np.column_stack)
+    jnp_fun = jnp.column_stack
+    self._CheckAgainstNumpy(jnp_fun, np_fun, args_maker)
+    self._CompileAndCheck(jnp_fun, args_maker)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_axis={}".format(
           jtu.format_test_name_suffix("", [shape] * len(dtypes), dtypes), axis),
Repository: google/jax

Issue:
custom_jvp does not work when first jvp rule is None
Here is a repro code
```python
from jax import grad, custom_jvp

@custom_jvp
def f(x, y):
  return x ** 2 * y

f.defjvps(None, lambda x_dot, primal_out, x, y: 2 * x * y * x_dot)
print(grad(f, 1)(2., 3.))
``` 


Discussion:

Resolving patch:
diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -75,7 +75,7 @@ def _initial_style_jaxpr(fun, in_avals):
   typed_jaxpr = core.TypedJaxpr(jaxpr, consts, in_avals, out_avals)
   return typed_jaxpr
 
-def sum_tangents(x, *xs):
+def sum_tangents(_, x, *xs):
   return reduce(ad.add_tangents, xs, x)
 
 def zeros_like_pytree(x):
@@ -196,7 +196,7 @@ def jvp(primals, tangents):
       zeros = zeros_like_pytree(primal_out)
       all_tangents_out = [jvp(t, primal_out, *primals) if jvp else zeros
                           for t, jvp in zip(tangents, jvps)]
-      tangent_out = tree_multimap(sum_tangents, *all_tangents_out)
+      tangent_out = tree_multimap(sum_tangents, primal_out, *all_tangents_out)
       return primal_out, tangent_out
 
     self.defjvp(jvp)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2476,6 +2476,17 @@ def foo(x):
     expected = -1.
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def test_custom_jvps_first_rule_is_none(self):
+    # https://github.com/google/jax/issues/3389
+    @api.custom_jvp
+    def f(x, y):
+      return x ** 2 * y
+
+    f.defjvps(None, lambda x_dot, primal_out, x, y: 2 * x * y * x_dot)
+    ans = grad(f, 1)(2., 3.)  # doesn't crash
+    expected = 12.
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
 
 class CustomVJPTest(jtu.JaxTestCase):
Repository: google/jax

Issue:
concatenate fails when axis = None
when running documentation example I'm getting a error:
```
import jax.numpy as np
a = np.array([[1, 2], [3, 4]])  
b = np.array([[5, 6]])  
np.concatenate((a, b), axis=None)  
```
Error:
```
---------------------------------------------------------------------------

TypeError                                 Traceback (most recent call last)

<ipython-input-37-28d5f29078d3> in <module>()
      5 np.concatenate((a, b.T), axis=1)
      6 
----> 7 np.concatenate((a, b), axis=None)
      8 

1 frames

/usr/local/lib/python3.6/dist-packages/jax/numpy/lax_numpy.py in concatenate(arrays, axis)
   1996   if ndim(arrays[0]) == 0:
   1997     raise ValueError("Zero-dimensional arrays cannot be concatenated.")
-> 1998   axis = _canonicalize_axis(axis, ndim(arrays[0]))
   1999   arrays = _promote_dtypes(*arrays)
   2000   # lax.concatenate can be slow to compile for wide concatenations, so form a

/usr/local/lib/python3.6/dist-packages/jax/lax/lax.py in _canonicalize_axis(axis, num_dims)
   5490 def _canonicalize_axis(axis, num_dims):
   5491   """Canonicalize an axis in [-num_dims, num_dims) to [0, num_dims)."""
-> 5492   axis = operator.index(axis)
   5493   if not -num_dims <= axis < num_dims:
   5494       raise ValueError(

TypeError: 'NoneType' object cannot be interpreted as an integer


```


Discussion:
JAX currently does not yet support the `axis=None` behavior of NumPy (flattening before concatenation). That said, the error message here is also quite poor.

Contributions would be welcome here -- either to implement the feature or a better error message! Take a look at `concatenate` in `jax/numpy/lax_numpy.py`.

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
@@ -1995,6 +1995,8 @@ def concatenate(arrays, axis=0):
     raise ValueError("Need at least one array to concatenate.")
   if ndim(arrays[0]) == 0:
     raise ValueError("Zero-dimensional arrays cannot be concatenated.")
+  if axis is None:
+    return concatenate([ravel(a) for a in arrays], axis=0)
   axis = _canonicalize_axis(axis, ndim(arrays[0]))
   arrays = _promote_dtypes(*arrays)
   # lax.concatenate can be slow to compile for wide concatenations, so form a


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
@@ -1218,6 +1218,12 @@ def args_maker():
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  def testConcatenateAxisNone(self):
+    # https://github.com/google/jax/issues/3419
+    a = jnp.array([[1, 2], [3, 4]])
+    b = jnp.array([[5]])
+    jnp.concatenate((a, b), axis=None)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_axis={}_baseshape=[{}]_dtypes=[{}]".format(
           axis, ",".join(str(d) for d in base_shape),
Repository: google/jax

Issue:
jax.numpy.arange() doesn't match numpy.arange() for fractional values
Example:
```python
In [1]: import jax.numpy as jnp    

In [2]: import numpy as np                                                                                         

In [3]: jnp.arange(2.5)                                                                                          
Out[3]: DeviceArray([0., 1.], dtype=float32)
                                                                            
In [4]: np.arange(2.5)                                                                                           
Out[4]: array([0., 1., 2.])
```


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
@@ -2253,7 +2253,7 @@ def arange(start, stop=None, step=None, dtype=None):
   lax._check_user_dtype_supported(dtype, "arange")
   if stop is None and step is None:
     dtype = dtype or _dtype(start)
-    return lax.iota(dtype, start) # avoids materializing
+    return lax.iota(dtype, ceil(start)) # avoids materializing
   else:
     return array(np.arange(start, stop=stop, step=step, dtype=dtype))
 


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
@@ -2504,9 +2504,11 @@ def testNpMean(self):
 
   def testArangeOnFloats(self):
     # from https://github.com/google/jax/issues/145
-    expected = np.arange(0.0, 1.0, 0.1, dtype=jnp.float_)
-    ans = jnp.arange(0.0, 1.0, 0.1)
-    self.assertAllClose(expected, ans)
+    self.assertAllClose(np.arange(0.0, 1.0, 0.1, dtype=jnp.float_),
+                        jnp.arange(0.0, 1.0, 0.1))
+    # from https://github.com/google/jax/issues/3450
+    self.assertAllClose(np.arange(2.5, dtype=jnp.float_),
+                        jnp.arange(2.5))
 
   def testSortManually(self):
     # manual tests for sort are nice because we don't have to worry about ties.
Repository: google/jax

Issue:
Module level documentation no longer shown on RTD
Example:
https://jax.readthedocs.io/en/latest/jax.experimental.loops.html

This used to show the module level docstring, but not anymore.

My guess is that this might be related to inserting the line `from __future__ import annotations` into each source file?


Discussion:

Resolving patch:
diff --git a/docs/conf.py b/docs/conf.py
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -28,30 +28,14 @@
 #
 import os
 import sys
-import typing
 
 sys.path.insert(0, os.path.abspath('..'))
 
 
-# Workaround to avoid expanding type aliases. See:
+# Currently type aliases are expanded. We tried a workaround along the lines of:
 # https://github.com/sphinx-doc/sphinx/issues/6518#issuecomment-589613836
-
-# When building docs, enable `from __future__ import annotations` everywhere.
-def _rewrite(p):
-  with open(p) as f:
-    contents = f.read()
-  with open(p, 'w') as f:
-    f.write('from __future__ import annotations\n')
-    f.write(contents)
-
-if 'READTHEDOCS' in os.environ:
-  for path, dirs, files in os.walk('../jax/'):
-    for file in files:
-      if file.endswith('.py'):
-        _rewrite(os.path.abspath(os.path.join(path, file)))
-
-# Monkey patch for the typing module to prevent it from expanding type aliases.
-typing.get_type_hints = lambda obj, *unused: obj.__annotations__
+# Unfortunately, this workaround makes Sphinx drop module-level documentation.
+# See https://github.com/google/jax/issues/3452.
 
 # -- Project information -----------------------------------------------------
 


Test patch:
Repository: google/jax

Issue:
fix numpy 1.19 deprecation warnings
Follow-up to #3509, the numpy 1.19 release started raising some DeprecationWarnings in tests that made our CI mad.


Discussion:

Resolving patch:
diff --git a/jax/dtypes.py b/jax/dtypes.py
--- a/jax/dtypes.py
+++ b/jax/dtypes.py
@@ -144,8 +144,10 @@ def _issubclass(a, b):
 
 def issubdtype(a, b):
   if a == bfloat16:
-    return b in [bfloat16, _bfloat16_dtype, np.floating, np.inexact,
-                 np.number]
+    if isinstance(b, np.dtype):
+      return b == _bfloat16_dtype
+    else:
+      return b in [bfloat16, np.floating, np.inexact, np.number]
   if not _issubclass(b, np.generic):
     # Workaround for JAX scalar types. NumPy's issubdtype has a backward
     # compatibility behavior for the second argument of issubdtype that
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
@@ -139,7 +139,7 @@ def __hash__(self):
     return hash(self.dtype.type)
 
   def __eq__(self, other):
-    return id(self) == id(other) or self.dtype == other
+    return id(self) == id(other) or self.dtype.type == other
 
   def __ne__(self, other):
     return not (self == other)
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -29,7 +29,7 @@
     packages=find_packages(exclude=["examples"]),
     python_requires='>=3.6',
     install_requires=[
-        'numpy >=1.12, <1.19', 'absl-py', 'opt_einsum'
+        'numpy >=1.12', 'absl-py', 'opt_einsum'
     ],
     url='https://github.com/google/jax',
     license='Apache-2.0',


Test patch:
Repository: google/jax

Issue:
Forward mode differentiation of `odeint` is not supported, but it does not give an error
```
def f(k):
  return ode.odeint(lambda x, t: k*x, 1.,  jnp.linspace(0, 1., 50)).sum()
jax.jacfwd(f)(3.)
```
should fail with an error since jvp of `odeint` is not supported, but it returns a value anyway.


Discussion:
Cross-ref to #3557 also: one reason this error isn't raised is that we don't detect the closed-over tracer in the dynamics function.
#2718 requests support for closure under `odeint`, and describes a related issue with how it is caught in the meantime.
In a funny way, this is accidentally doing the right thing, even though I told @john-m-jumper it wouldn't!

```python
import jax
from jax.experimental import ode
import jax.numpy as jnp

def f(k):
  return ode.odeint(lambda x, t: k*x, 1.,  jnp.linspace(0, 1., 50)).sum()
print(jax.jacfwd(f)(3.))
print((f(3.+1e-6) - f(3.-1e-6)) / 2e-6)
```

```
$ env JAX_ENABLE_X64=1 python issue3558.py
234.33132392618268
234.33132375316745
```

It's kind of funny: `odeint` has a custom vjp rule, and usually when one has a custom vjp rule that means we can't use forward-mode with it. However, `odeint` also isn't set up to detect closed-over tracers in the dynamics function. As a result, the tracer on `k` is eluding odeint's custom vjp rule, but that means it just gets into odeint's implementation, which is all implemented in terms of while_loops and things that are themselves forward-mode differentiable. So it all just works out!

We should not close this bug yet though; we want to make sure this keeps working when we resolve #3557 / #2718 , and moreover make sure forward-mode works like this even without smuggling a tracer in via closure.
After merging #3562, now we get an error here:

```
NotImplementedError: Batching rule for 'custom_lin' not implemented
```

That's a bad error message, but it makes sense: it's essentially from doing forward-mode on a function with a custom VJP.

At least two possible directions to go:
1. make forward-mode work with custom vjp functions, perhaps ignoring the custom rule and just tracing through the function; or
2. improve the error message.

Resolving patch:
diff --git a/jax/custom_derivatives.py b/jax/custom_derivatives.py
--- a/jax/custom_derivatives.py
+++ b/jax/custom_derivatives.py
@@ -590,3 +590,5 @@ def batched_fwd_jaxpr_thunk():
 
 xla.initial_style_translations[custom_vjp_call_jaxpr_p] = \
     xla.lower_fun_initial_style(_custom_vjp_call_jaxpr_impl)
+
+batching.primitive_batchers[ad.custom_lin_p] = ad._raise_custom_vjp_error_on_jvp


Test patch:
diff --git a/tests/ode_test.py b/tests/ode_test.py
--- a/tests/ode_test.py
+++ b/tests/ode_test.py
@@ -212,6 +212,15 @@ def model(y, t):
 
     self.assertAllClose(ans, expected, check_dtypes=False, atol=1e-2, rtol=1e-2)
 
+  def test_forward_mode_error(self):
+    # https://github.com/google/jax/issues/3558
+
+    def f(k):
+      return odeint(lambda x, t: k*x, 1.,  jnp.linspace(0, 1., 50)).sum()
+
+    with self.assertRaisesRegex(TypeError, "can't apply forward-mode.*"):
+      jax.jacfwd(f)(3.)
+
 
 if __name__ == '__main__':
   absltest.main()
Repository: google/jax

Issue:
lax.pad breaks for zero-sized inputs
```python
from jax import lax, numpy as jnp
out = lax.pad(jnp.ones((0,)), 0., ((1, 1, 1),))

print(out.shape) # (1,)
print(out)       # [0. 0.]
print(out[0])    # RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 0: want f32[1]{0}, got f32[2]{0}
```

Pad works as expected for zero-sized inputs for non-interior padding (i. e. padding config `((1, 1, 0),)`), so I guess this should also work (or at least give an error).


Discussion:

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -2950,8 +2950,9 @@ def _pad_dtype_rule(operand, padding_value, *, padding_config):
 
 def _pad_shape_rule(operand, padding_value, *, padding_config):
   lo, hi, interior = zip(*padding_config)
-  out_shape = onp.add(onp.add(onp.add(lo, hi), operand.shape),
-                      onp.multiply(interior, onp.subtract(operand.shape, 1)))
+  out_shape = onp.add(
+    onp.add(onp.add(lo, hi), operand.shape),
+    onp.maximum(0, onp.multiply(interior, onp.subtract(operand.shape, 1))))
   return tuple(out_shape)
 
 def _pad_transpose(t, operand, padding_value, *, padding_config):


Test patch:
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -955,7 +955,7 @@ def testReshapeAgainstNumpy(self, arg_shape, out_shape, dtype, rng_factory):
       {"testcase_name": "_inshape={}_pads={}"
        .format(jtu.format_shape_dtype_string(shape, dtype), pads),
        "shape": shape, "dtype": dtype, "pads": pads, "rng_factory": jtu.rand_small}
-      for shape in [(2, 3)]
+      for shape in [(0, 2), (2, 3)]
       for dtype in default_dtypes
       for pads in [[(1, 2, 1), (0, 1, 0)]]))
   def testPad(self, shape, dtype, pads, rng_factory):
Repository: google/jax

Issue:
broken link from pytree_util api docs to 'pytrees notebook'
the bottom of the top introduction on this page:
https://jax.readthedocs.io/en/latest/jax.tree_util.html

reads: 

> See the [JAX pytrees notebook](https://jax.readthedocs.io/en/latest/notebooks/JAX_pytrees.html) for examples.

(this is written here: https://github.com/google/jax/blob/269da0ae584cfe840f34e9f871f13c28e2772de5/jax/tree_util.py#L34)

is this the correct link? or are we missing a notebook?
https://jax.readthedocs.io/en/latest/pytrees.html

just curious because tree_util is my new favorite thing...


Discussion:

Resolving patch:
diff --git a/jax/tree_util.py b/jax/tree_util.py
--- a/jax/tree_util.py
+++ b/jax/tree_util.py
@@ -31,7 +31,7 @@
 user defined data structures and JAX transformations (e.g. `jit`). This is not
 meant to be a general purpose tree-like data structure handling library.
 
-See the `JAX pytrees notebook <https://jax.readthedocs.io/en/latest/notebooks/JAX_pytrees.html>`_
+See the `JAX pytrees note <pytrees.html>`_
 for examples.
 """
 
@@ -112,7 +112,7 @@ def all_leaves(iterable):
 def register_pytree_node(nodetype, flatten_func, unflatten_func):
   """Extends the set of types that are considered internal nodes in pytrees.
 
-  See `example usage <https://jax.readthedocs.io/en/latest/notebooks/JAX_pytrees.html#Pytrees-are-extensible>`_.
+  See `example usage <pytrees.html>`_.
 
   Args:
     nodetype: a Python type to treat as an internal pytree node.


Test patch: