tuple_replace function in util.py seems incorrect
In the file jax._src.util.py

Current definition
```python
def tuple_replace(t, idx, val):
  assert 0 <= idx < len(t), (idx, len(t))
  return t[:idx] + (val,) + t[idx:]
```
looks like an almost copy of tuple_insert. I think the correct definition should be:

```python
def tuple_replace(t, idx, val):
  assert 0 <= idx < len(t), (idx, len(t))
  return t[:idx] + (val,) + t[idx+1:]
```


diff --git a/jax/_src/util.py b/jax/_src/util.py
index 607479216..051f570f3 100644
--- a/jax/_src/util.py
+++ b/jax/_src/util.py
@@ -330,10 +330,6 @@ def tuple_delete(t, idx):
   assert 0 <= idx < len(t), (idx, len(t))
   return t[:idx] + t[idx + 1:]
 
-def tuple_replace(t, idx, val):
-  assert 0 <= idx < len(t), (idx, len(t))
-  return t[:idx] + (val,) + t[idx:]
-
 # TODO(mattjj): replace with dataclass when Python 2 support is removed
 def taggedtuple(name, fields) -> Callable[..., Any]:
   """Lightweight version of namedtuple where equality depends on the type."""
['jax/_src/util.py']