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

Issue:
Implement `np.select`
It appears that JAX has not implemented `np.select`.

One can work around this by using `lax.select`, but it doesn't implement numpy's broadcasting semantics, so it's not quite as nice.


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
@@ -768,6 +768,20 @@ def where(condition, x=None, y=None):
     return lax.select(condition, *_promote_dtypes(x, y))
 
 
+@_wraps(onp.select)
+def select(condlist, choicelist, default=0):
+  if len(condlist) != len(choicelist):
+    msg = "condlist must have length equal to choicelist ({} vs {})"
+    raise ValueError(msg.format(len(condlist), len(choicelist)))
+  if len(condlist) == 0:
+    raise ValueError("condlist must be non-empty")
+
+  output = default
+  for cond, choice in zip(condlist[::-1], choicelist[::-1]):
+    output = where(cond, choice, output)
+  return output
+
+
 def broadcast_arrays(*args):
   """Like Numpy's broadcast_arrays but doesn't return views."""
   shapes = [shape(arg) for arg in args]


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
@@ -1478,6 +1478,29 @@ def testIx_(self, rng, shapes, dtypes):
                             check_dtypes=True)
     self._CompileAndCheck(lnp.ix_, args_maker, check_dtypes=True)
 
+
+  @parameterized.named_parameters(jtu.cases_from_list(
+        {"testcase_name": jtu.format_test_name_suffix("select", shapes,
+                                                      (onp.bool_,) * n + dtypes),
+         "rng": jtu.rand_default(), "shapes": shapes, "dtypes": dtypes}
+        for n in range(0, 3)
+        for shapes in filter(
+          _shapes_are_broadcast_compatible,
+          CombosWithReplacement(all_shapes, 2 * n + 1))
+        for dtypes in CombosWithReplacement(all_dtypes, n + 1)))
+  def test(self, rng, shapes, dtypes):
+    n = len(dtypes) - 1
+    def args_maker():
+      condlist = [rng(shape, onp.bool_) for shape in shapes[:n]]
+      choicelist = [rng(shape, dtype)
+                    for shape, dtype in zip(shapes[n:-1], dtypes[:n])]
+      default = rng(shapes[-1], dtypes[-1])
+      return condlist, choicelist, default
+    self._CheckAgainstNumpy(onp.select, lnp.select, args_maker,
+                            check_dtypes=True)
+    self._CompileAndCheck(lnp.select, args_maker, check_dtypes=True)
+
+
   def testIssue330(self):
     x = lnp.full((1, 1), lnp.array([1])[0])  # doesn't crash
     self.assertEqual(x[0, 0], 1)
Repository: google/jax

Issue:
Large memory needed in gathering operations
[Here][1] is a simplified example of my problem. While I can compute the gradients of a single input, I quickly run out of memory. Is there a way to make it more memory efficient?

With a batch size of 6 I need 2157 MiB, but the gradients are of shapes (6, 50, 21) and (6, 50, 50, 21, 21), which are 25 MiB, but I am incurring in a 100x overhead. Am I doing something wrong? Can we help jax be smarter? Can I profile or inspect the graph to figure out which operations are taking the space?



[1]: https://gist.github.com/Dapid/e980a26e063838d472f46db1b7df59bd


Discussion:
From a quick look, this looks very similar to #824.

Can you try to use a single large gather (e.g., using advanced indexing) rather multiple lookups in a loop?

(That said, I suspect we can improve the memory usage of your code as written, too.)
The problem of #824 was dynamic slice. I don't think I can get rid of the loop, in the real case each iteration depends on `r`, but I will rethink it.
Nice code!

Under a `jit` I would expect the sparse gradients to be accumulated efficiency, but we should investigate if something is going wrong there. In Autograd we also accumulated sparse gradients that arose from indexing efficiently (without having `jit`) because Dougal added [a special SparseObject](https://github.com/HIPS/autograd/blob/541a70b8a754356fe52e2d6caa7f4aa94c18cf05/autograd/numpy/numpy_vjps.py#L700) that [the gradient-summing machinery knew how to deal with](https://github.com/HIPS/autograd/blob/541a70b8a754356fe52e2d6caa7f4aa94c18cf05/autograd/core.py#L158-L186). We could do that in JAX too if necessary, though I'd be interested to figure out why `jit` might not be doing the optimization here.

Resolving patch:
diff --git a/jax/lax/lax.py b/jax/lax/lax.py
--- a/jax/lax/lax.py
+++ b/jax/lax/lax.py
@@ -2630,7 +2630,7 @@ def _dynamic_slice_jvp_rule(g, operand, start_indices, slice_sizes,
 def _dynamic_slice_transpose_rule(t, operand, start_indices, slice_sizes,
                                   operand_shape):
   assert operand is None
-  zeros = full(operand_shape, 0, dtype=_dtype(t))
+  zeros = full(operand_shape, tie_in(t, _zero(t)))
   return [dynamic_update_slice(zeros, t, start_indices), ad_util.zero]
 
 def _dynamic_slice_batching_rule(batched_args, batch_dims, slice_sizes,
@@ -2798,7 +2798,7 @@ def _gather_transpose_rule(t, operand, start_indices, dimension_numbers,
   assert operand is None
   if t is ad_util.zero:
     return [ad_util.zero, ad_util.zero]
-  zeros = full(operand_shape, 0, dtype=t.dtype)
+  zeros = full(operand_shape, tie_in(t, _zero(t)))
   scatter_dnums = ScatterDimensionNumbers(
     update_window_dims=dimension_numbers.offset_dims,
     inserted_window_dims=dimension_numbers.collapsed_slice_dims,


Test patch:
Repository: google/jax

Issue:
`jax.device_put` has complexity superlinear in the number of arguments.
Reproduction on CPU:
```
import jax
import numpy as onp

In [9]: %time x = jax.device_put([onp.random.randn(10,5) for _ in range(100)])
CPU times: user 1.45 s, sys: 7.8 ms, total: 1.46 s
Wall time: 1.45 s

In [10]: %time x = jax.device_put([onp.random.randn(10,5) for _ in range(500)])
CPU times: user 24.8 s, sys: 0 ns, total: 24.8 s
Wall time: 24.8 s

In [11]: %time x = jax.device_put([onp.random.randn(10,5) for _ in range(700)])
CPU times: user 45.2 s, sys: 190 ms, total: 45.4 s
Wall time: 45.4 s
```


Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -989,7 +989,26 @@ def jaxpr_maker(*args, **kwargs):
 tree_to_pval_tuples = partial(process_pytree, pe.pack_pvals)
 
 
-device_put = jit(lambda x: x)
+
+_traceable_device_put = jit(lambda x: x)
+
+def device_put(x, device_num=0):
+  def _device_put(x):
+    if isinstance(x, core.Tracer):
+      return _traceable_device_put(x)
+
+    try:
+      a = xla.abstractify(x)
+    except TypeError:
+      raise TypeError("Argument '{}' of type {} is not a valid JAX type"
+                      .format(x, type(x)))
+
+    result_shape = xla.xla_shape_to_result_shape(xla.xla_shape(a))
+    handler = xla.device_persistent_result_handler(result_shape)
+    return handler(xla.device_put(x, device_num))
+  return tree_map(_device_put, x)
+
+
 device_get = _jit(lambda x: x, (), device_values=False)
 
 
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -179,11 +179,11 @@ class _ResultArray(tuple): pass
 
 def result_handler(result_shape):
   if FLAGS.jax_device_values:
-    return _device_persistent_result_handler(result_shape)
+    return device_persistent_result_handler(result_shape)
   else:
     return _pyval_result_handler(result_shape)
 
-def _device_persistent_result_handler(result_shape):
+def device_persistent_result_handler(result_shape):
   t = type(result_shape)
   if t is _ResultArray:
     return partial(DeviceArray, result_shape)
@@ -416,7 +416,7 @@ def __init__(self, result_shape, device_buffer):
 
   def __iter__(self):
     bufs = self.device_buffer.destructure()
-    handlers = map(_device_persistent_result_handler, self.result_shapes)
+    handlers = map(device_persistent_result_handler, self.result_shapes)
     elts = [handler(buf) for handler, buf in zip(handlers, bufs)]
     return iter(elts)
 
@@ -632,7 +632,7 @@ def _xla_callable(fun, device_values, *abstract_args):
     compiled, result_shape = _compile_jaxpr(jaxpr, consts, *abstract_args)
     del master, consts, jaxpr, env
   if device_values:
-    handle_result = _device_persistent_result_handler(result_shape)
+    handle_result = device_persistent_result_handler(result_shape)
   else:
     handle_result = _pyval_result_handler(result_shape)
   return partial(_execute_compiled, compiled, pval, handle_result)


Test patch:
Repository: google/jax

Issue:
`list(np.arange(10001))` is slow.
We should improve this.


Discussion:

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -526,13 +526,13 @@ def __iter__(self):
     if self.ndim == 0:
       raise TypeError("iteration over a 0-d array")  # same as numpy error
     else:
-      return (self[i] for i in xrange(self.shape[0]))
+      return self._value.__iter__()
 
   def __reversed__(self):
     if self.ndim == 0:
       raise TypeError("iteration over a 0-d array")
     else:
-      return (self[i] for i in xrange(self.shape[0] - 1, -1, -1))
+      return reversed(self._value)
 
   def __format__(self, format_spec):
     # Simulates behavior of https://github.com/numpy/numpy/pull/9883


Test patch:
Repository: google/jax

Issue:
convert integer for `static_argnums` to tuple automatically
I've made the same mistake a few times on `static_argnums` on `jit` function, by not using a tuple 

e.g., `jit(f, static_argnums=1)`
would it be possible / reasonable to add a simple `if` statement to convert it to tuple?

where the correct way is  `jit(f, static_argnums=(1,) )`

also note `(1)` is not a tuple,  but `(1,)` is




Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -110,6 +110,9 @@ def _jit(fun, static_argnums, device_values=True):
   def f_jitted(*args, **kwargs):
     if _jit_is_disabled or config.read('jax_disable_jit'):
       return fun(*args, **kwargs)
+    if isinstance(static_argnums, int):
+      # static_argnums is a tuple of ints
+      static_argnums = (static_argnums,)
     if static_argnums and max(static_argnums) >= len(args):
       msg = ("Jitted function has static_argnums={} but was called with only {}"
              " positional arguments.")


Test patch:
Repository: google/jax

Issue:
[FR] support batching rule for np.cumprod
It is strange to me that we have batching rule for `reduce_window` but `np.cumprod` raises `NotImplementedError: Batching rule for 'reduce_window' not implemented`. Here is a simple repro code
```python
from jax import vmap
import jax.numpy as np

def f(x):
    return np.cumprod(x, axis=-1)

vmap(f)(np.ones((3, 3)))
```


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
@@ -3397,9 +3397,27 @@ def _reduce_window_translation_rule(c, operand, init_value, jaxpr, consts,
   return c.ReduceWindow(operand, init_value, xla_computation, window_dimensions,
                         window_strides, padding)
 
+def _generic_reduce_window_batch_rule(
+    batched_args, batch_dims, jaxpr, consts, window_dimensions, window_strides,
+    padding):
+  operand, init = batched_args
+  bdim, init_bdim = batch_dims
+  if init_bdim is not None:
+    raise NotImplementedError("reduce_window batching is not implemented for "
+                              "initial values")
+
+  def reduce_window(x, window_dimensions, window_strides, padding):
+    return reduce_window_p.bind(
+      x, init, jaxpr=jaxpr, consts=consts, window_dimensions=window_dimensions,
+      window_strides=window_strides, padding=padding)
+  return _reduce_window_batch_rule(reduce_window, (operand,), (bdim,),
+                                   window_dimensions, window_strides, padding)
+
+
 reduce_window_p = standard_primitive(
     _reduce_window_shape_rule, _input_dtype, 'reduce_window',
     _reduce_window_translation_rule)
+batching.primitive_batchers[reduce_window_p] = _generic_reduce_window_batch_rule
 
 
 def _reduce_window_sum_shape_rule(operand, window_dimensions, window_strides,


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -548,6 +548,11 @@ def f(params, x):
     per_example_direct = np.concatenate(per_example_direct, axis=0)
     self.assertAllClose(per_example, per_example_direct, check_dtypes=True)
 
+  def testCumProd(self):
+   x = np.arange(9).reshape(3, 3) + 1
+   y = vmap(lambda x: np.cumprod(x, axis=-1))(x)
+   self.assertAllClose(onp.cumprod(x, axis=1), y, check_dtypes=True)
+
   def testSelect(self):
     pred = onp.array([True, False])
     on_true = onp.array([0, 1])
Repository: google/jax

Issue:
`jax.jit` should report an error immediately for uncallable arguments
`jax.jit(3)` silently succeeds (although `jax.jit(3)(4)` returns the expected error). The `callable` builtin seems applicable here.



Discussion:

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- a/jax/api.py
+++ b/jax/api.py
@@ -68,6 +68,10 @@
                   "Disable JIT compilation and just call original Python.")
 
 
+def _check_callable(fun):
+  if not callable(fun):
+    raise TypeError("Expected a callable value, got {}".format(fun))
+
 def jit(fun, static_argnums=()):
   """Sets up `fun` for just-in-time compilation with XLA.
 
@@ -106,6 +110,7 @@ def jit(fun, static_argnums=()):
   return _jit(fun, static_argnums)
 
 def _jit(fun, static_argnums, device_values=True):
+  _check_callable(fun)
   if isinstance(static_argnums, int):
     static_argnums = (static_argnums,)
 
@@ -248,6 +253,7 @@ def xla_computation(fun, static_argnums=(), axis_env=None):
     ROOT tuple.18 = (f32[], f32[], f32[]) tuple(all-reduce.7, all-reduce.12, all-reduce.17)
   }
   """
+  _check_callable(fun)
 
   def pv_like(x):
     aval = xla.abstractify(x)
@@ -362,6 +368,8 @@ def value_and_grad(fun, argnums=0, has_aux=False, holomorphic=False):
             "of {fun} and the second element is the gradient, which has the "
             "same shape as the arguments at positions {argnums}.")
 
+  _check_callable(fun)
+
   @wraps(fun, docstr=docstr, argnums=argnums)
   def value_and_grad_f(*args, **kwargs):
     f = lu.wrap_init(fun, kwargs)
@@ -568,6 +576,8 @@ def vmap(fun, in_axes=0, out_axes=0):
   docstr = ("Vectorized version of {fun}. Takes similar arguments as {fun} "
             "but with additional array axes over which {fun} is mapped.")
 
+  _check_callable(fun)
+
   if (not isinstance(in_axes, (list, tuple, type(None), int))
       or not isinstance(out_axes, (list, tuple, type(None), int))):
     msg = ("vmap arguments in_axes and out_axes must each be an integer, None, "
@@ -674,6 +684,7 @@ def pmap(fun, axis_name=None):
   >>> print(doubly_normed.sum((0, 1)))
   1.0
   """
+  _check_callable(fun)
   axis_name = _TempAxisName() if axis_name is None else axis_name
 
   @wraps(fun)
@@ -723,6 +734,7 @@ def _aval_axis_size(aval):
 
 
 def soft_pmap(fun, axis_name=None):
+  _check_callable(fun)
   axis_name = _TempAxisName() if axis_name is None else axis_name
 
   @wraps(fun)
@@ -1066,6 +1078,8 @@ def make_jaxpr(fun):
         (l) = id k
     in l }
   """
+  _check_callable(fun)
+
   def pv_like(x):
     aval = xla.abstractify(x)
     return pe.PartialVal((aval, core.unit))


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -895,6 +895,9 @@ def f(x):
     xla_comp = api.xla_computation(f)
     xla_comp(np.arange(8)).GetHloText()  # doesn't crash
 
+  def test_jit_of_noncallable(self):
+    jtu.check_raises_regexp(lambda: api.jit(3), TypeError,
+                            "Expected a callable value.*")
 
 if __name__ == '__main__':
   absltest.main()
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.


Discussion:
I think the underlying issue is that [we only look at `args` to determine the axis size](https://github.com/google/jax/blob/39b01564d4c6776960a45898bdcd8b17f517469a/jax/api.py#L714), even though once the axis size is determined we handle kwargs correctly.

Resolving patch:
diff --git a/jax/api.py b/jax/api.py
--- 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)


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -955,6 +955,14 @@ def testIndexAddBatchedIndexesOnly(self):
     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.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.
```


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
@@ -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.")


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
@@ -1864,6 +1864,18 @@ def testDisableNumpyRankPromotionBroadcasting(self):
     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:
block_until_ready() should return the DeviceArray object
It would be nice to have the `block_until_ready()` method return the `DeviceArray` object so that it can be more easily dropped into existing code.

So something like 

```python
a = f()
```

could easily become

```python
a = f().block_until_ready()
```


Discussion:

Resolving patch:
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
--- a/jax/interpreters/pxla.py
+++ b/jax/interpreters/pxla.py
@@ -297,6 +297,7 @@ def block_until_ready(self):
     self._check_if_deleted()
     for buf in self.device_buffers:
       buf.block_host_until_ready()
+    return self
 
 
 class ShardedDeviceArray(ShardedDeviceValue, xla.DeviceArray):
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -507,9 +507,12 @@ def block_until_ready(self):
     This method is mostly useful for timing microbenchmarks that wish to
     time how long a computation takes, without transferring the result back
     to the host.
+
+    Returns the buffer object (`self`).
     """
     self._check_if_deleted()
     self.device_buffer.block_host_until_ready()
+    return self
 
 def _forward_method(attrname, self, fun, *args):
   return fun(getattr(self, attrname), *args)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -699,8 +699,9 @@ def test_devicearray_delete(self):
 
   def test_devicearray_block_until_ready(self):
     x = device_put(1.)
-    x.block_until_ready()
-    # Tests only that block_until_ready() does not produce an error.
+    y = x.block_until_ready()
+    # Tests mostly that block_until_ready() does not produce an error.
+    self.assertTrue(y is x)
 
   def test_namedtuple_transparency(self):
     # See https://github.com/google/jax/issues/446
Repository: google/jax

Issue:
argmax instantiates global constants
The culprit is this line: https://github.com/google/jax/blob/22afc8cea1b9b5a0dca0cdfe8a0ce51df94cf9e0/jax/numpy/lax_numpy.py#L2234

I'm attempting to do argmax on large tensors and am running out of memory due to the size of the globally-allocated range object.

A `tie_in` might workaround might address the memory issue, but perhaps the better approach is to implement what the TODO says?
https://github.com/google/jax/blob/22afc8cea1b9b5a0dca0cdfe8a0ce51df94cf9e0/jax/numpy/lax_numpy.py#L2230-L2237


Discussion:
My current workaround is:

```
def _argminmax(op, a, axis):
  shape = [1] * a.ndim
  shape[axis] = a.shape[axis]
  idxs = lax.tie_in(a, arange(a.shape[axis])).reshape(shape)
  maxval = onp.iinfo(xla_bridge.canonicalize_dtype(idxs.dtype)).max
  maxval = lax.tie_in(a, maxval)
  mask_idxs = where(lax._eq_meet(a, op(a, axis, keepdims=True)), idxs, maxval)
  return min(mask_idxs, axis)
```

I realized that maxval needed a tie-in too, because it would be broadcast as part of `where`, with the result constant-folded and stored in global memory.

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
@@ -2232,8 +2232,9 @@ def argmin(a, axis=None):
 def _argminmax(op, a, axis):
   shape = [1] * a.ndim
   shape[axis] = a.shape[axis]
-  idxs = onp.arange(a.shape[axis]).reshape(shape)
+  idxs = lax.tie_in(a, arange(a.shape[axis])).reshape(shape)
   maxval = onp.iinfo(xla_bridge.canonicalize_dtype(idxs.dtype)).max
+  maxval = lax.tie_in(a, maxval)
   mask_idxs = where(lax._eq_meet(a, op(a, axis, keepdims=True)), idxs, maxval)
   return min(mask_idxs, axis)
 


Test patch: