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

Issue:
[sparse] BCOO._dedupe moves padded indices to the front
Example:
```python
import jax.numpy as jnp
from jax.experimental import sparse

data = jnp.array([1, 1])
indices = jnp.array([1, 1])[:, None]
shape = (2,)

x = sparse.BCOO((data, indices), shape=shape)
x = x._dedupe()
print(x.data, x.indices.ravel())
# [2 0] [1 0]

x = x._dedupe()
print(x.data, x.indices.ravel())
# [0 2] [0 1]
```
The second `_dedupe()` call should leave the representation unchanged.


Discussion:

Resolving patch:
diff --git a/jax/experimental/sparse/bcoo.py b/jax/experimental/sparse/bcoo.py
--- a/jax/experimental/sparse/bcoo.py
+++ b/jax/experimental/sparse/bcoo.py
@@ -67,17 +67,20 @@ def _dedupe_bcoo_one(data, indices):
   if indices.shape[1] == 0:
     return data, indices
 
-  # This is a fixed-size version of jnp.unique() with return_indices=True
-  # unique values are zero-filled at the end.
-  perm = jnp.lexsort(indices.T[::-1])
+  # The following is similar to
+  #   indices_unique, inv_idx = jnp.unique(indices, axis=0, return_inverse=True,
+  #                                        size=indices.shape[0], fill_value=0)
+  # but modified to keep padding at the end of the resulting arrays.
+  is_padding = (indices == 0).all(1) & (data == 0)
+  perm = jnp.lexsort(indices[:, ::-1].T)
   aux = indices[perm]
   mask = jnp.ones(indices.shape[0], dtype=bool)
   mask = mask.at[1:].set(jnp.any(aux[1:] != aux[:-1], 1))
+  mask = mask & ~is_padding[perm]  # this is the padding modification.
   imask = jnp.cumsum(mask) - 1
   indices_unique = jnp.where(mask[:, None], aux, 0)[jnp.argsort(~mask)]
   inv_idx = jnp.zeros_like(imask).at[perm].set(imask)
 
-  # With the above, de-duping is easy.
   data_unique = jnp.zeros_like(data).at[inv_idx].add(data)
   return data_unique, indices_unique
 


Test patch:
diff --git a/tests/sparse_test.py b/tests/sparse_test.py
--- a/tests/sparse_test.py
+++ b/tests/sparse_test.py
@@ -1059,6 +1059,16 @@ def test_bcoo_dedupe(self, shape, dtype, n_batch, n_dense):
     M_dedup = M._dedupe()
     self.assertAllClose(M.todense(), M_dedup.todense())
 
+  def test_bcoo_dedupe_padding(self):
+    # Regression test for https://github.com/google/jax/issues/8163
+    data = jnp.array([1, 0, 0])
+    indices = jnp.array([1, 0, 0])[:, None]
+    x = sparse.BCOO((data, indices), shape=(3,))
+    y = x._dedupe()
+    self.assertArraysEqual(x.todense(), y.todense())
+    self.assertArraysEqual(x.indices, y.indices)
+    self.assertArraysEqual(x.data, y.data)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_{}_nbatch={}_ndense={}_axes={}".format(
         jtu.format_shape_dtype_string(shape, dtype), n_batch, n_dense, axes),
Repository: google/jax

Issue:
vmap documentation does not document "axis_name" parameter
vmap() accepts an axis_name parameter (as verified by looking at the code), but this parameter is not documented at all in the docstring.

(From visiting https://jax.readthedocs.io/en/latest/_modules/jax/_src/api.html#vmap just prior to filing this bug)




Discussion:
Thanks for the report - this was added in https://github.com/google/jax/pull/4005. I'm going to assign @apaszke as the author of that PR.
I'm stealing this!

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -1329,6 +1329,8 @@ def vmap(fun: F, in_axes=0, out_axes=0, axis_name=None) -> F:
       (axes) of the array returned by the :func:`vmap`-ed function, which is one
       more than the number of dimensions (axes) of the corresponding array
       returned by ``fun``.
+    axis_name: Optional, a hashable Python object used to identify the mapped
+      axis so that parallel collectives can be applied.
 
   Returns:
     Batched/vectorized version of ``fun`` with arguments that correspond to
@@ -1403,6 +1405,16 @@ def vmap(fun: F, in_axes=0, out_axes=0, axis_name=None) -> F:
 
   If the ``out_axes`` is specified for a mapped result, the result is transposed
   accordingly.
+
+  Finally, here's an example using ``axis_name`` together with collectives:
+
+  >>> xs = jnp.arange(3. * 4.).reshape(3, 4)
+  >>> print(vmap(lambda x: lax.psum(x, 'i'), axis_name='i')(xs))
+  [[12. 15. 18. 21.]
+   [12. 15. 18. 21.]
+   [12. 15. 18. 21.]]
+
+  See the :py:func:`jax.pmap` docstring for more examples involving collectives.
   """
   _check_callable(fun)
   docstr = ("Vectorized version of {fun}. Takes similar arguments as {fun} "


Test patch:
Repository: google/jax

Issue:
[sparse] sparsify transform does not respect weak types
```python
from jax import config
config.update('jax_enable_x64', True)

import jax.numpy as jnp
from jax.experimental import sparse

M = jnp.arange(12, dtype='int32').reshape(3, 4)
Msp = sparse.BCOO.fromdense(M)

print((2 * M).dtype)  # int32
print((2 * Msp).dtype)  # int64
```
Found in the process of #8180


Discussion:

Resolving patch:
diff --git a/jax/experimental/sparse/transform.py b/jax/experimental/sparse/transform.py
--- a/jax/experimental/sparse/transform.py
+++ b/jax/experimental/sparse/transform.py
@@ -55,6 +55,7 @@
 from jax import core
 from jax import lax
 from jax import linear_util as lu
+import jax.numpy as jnp
 from jax._src.api_util import flatten_fun_nokwargs
 from jax.interpreters import partial_eval as pe
 from jax.interpreters import xla
@@ -78,7 +79,7 @@ def __init__(self, bufs=()):
     self._buffers = list(bufs)
 
   def push(self, arr: Array) -> int:
-    self._buffers.append(np.array(arr) if np.isscalar(arr) else arr)  # type: ignore
+    self._buffers.append(jnp.asarray(arr))  # type: ignore
     return len(self._buffers) - 1
 
   def get(self, ind: int) -> Array:
@@ -155,7 +156,8 @@ def argspec_to_aval(argspec):
     if argspec.is_unit():
       return core.abstract_unit
     else:
-      return core.ShapedArray(argspec.shape, argspec.data(spenv).dtype)
+      data = argspec.data(spenv)
+      return core.ShapedArray(argspec.shape, data.dtype, data.aval.weak_type)
   return tree_map(argspec_to_aval, argspecs, is_leaf=_is_argspec)
 
 


Test patch:
diff --git a/tests/sparsify_test.py b/tests/sparsify_test.py
--- a/tests/sparsify_test.py
+++ b/tests/sparsify_test.py
@@ -343,6 +343,16 @@ def func(x, y):
     with self.assertRaisesRegex(TypeError, "sparsified true_fun and false_fun output.*"):
       func(x_bcoo, y)
 
+  def testWeakTypes(self):
+    # Regression test for https://github.com/google/jax/issues/8267
+    M = jnp.arange(12, dtype='int32').reshape(3, 4)
+    Msp = BCOO.fromdense(M)
+    self.assertArraysEqual(
+      operator.mul(2, M),
+      sparsify(operator.mul)(2, Msp).todense(),
+      check_dtypes=True,
+    )
+
 
 if __name__ == "__main__":
   absltest.main(testLoader=jtu.JaxTestLoader())
Repository: google/jax

Issue:
two-operand cond conflicts with per-branch-operand cond
#8467 introduces multi-operand `lax.cond`. When given two operands, as in:

```python
def f(x):
  return lax.cond(x == 0, lax.add, lax.mul, x, y)
```

`cond` detects five total arguments and falls back to behavior that is backwards-compatible with the deprecated per-branch-operand `cond`. The latter accepts, in order: a predicate, true operand, true branch, false operand, and false branch. This results in an error:

```
TypeError: Value <function add ...> with type <class 'function'> is not a valid JAX type
```

One option is to rollback (or disallow) multi-operand cond for now, complete the deprecation of the previous cond, and roll forward (or re-enable) multi-operand cond.


Discussion:
Having the same issue here! For now, I have to strictly use single operand `cond` calls.

Resolving patch:
diff --git a/jax/_src/lax/control_flow.py b/jax/_src/lax/control_flow.py
--- a/jax/_src/lax/control_flow.py
+++ b/jax/_src/lax/control_flow.py
@@ -788,7 +788,10 @@ def cond(*args, **kwargs):
   except TypeError:
     pass
   else:
-    return _cond_with_per_branch_args(*ba.args)
+    assert not ba.kwargs  # no catch-all **kwargs in _cond_with_per_branch
+    _, _, maybe_true_fun, _, maybe_false_fun = ba.args
+    if callable(maybe_true_fun) and callable(maybe_false_fun):
+      return _cond_with_per_branch_args(*ba.args)
 
   return _cond(*args, **kwargs)
 


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -599,8 +599,6 @@ def false_fun(x):
 
   def testCondTwoOperands(self):
     # see https://github.com/google/jax/issues/8469
-    self.skipTest("two-operand cond behavior is ambiguous (#8469)")
-
     add, mul = lax.add, lax.mul
 
     def fun(x):
Repository: google/jax

Issue:
jax.median fails for complex values
```python
>>> import numpy as np
>>> import jax.numpy as jnp
>>> x = np.array([1 + 0j, 0 + 1j])
>>> np.median(x)
(0.5+0.5j)
>>> jnp.median(x)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)

<...>

TypeError: floor does not accept dtype complex64. Accepted dtypes are subtypes of floating.
```


Discussion:
See https://github.com/numpy/numpy/issues/12943 for a discussion of the meaning of numpy's median for complex inputs; we may just want to raise an explicit `ValueError` or `NotImplementedError` rather than trying to replicate this.

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
@@ -6321,6 +6321,8 @@ def _quantile(a, q, axis, interpolation, keepdims, squash_nans):
     raise ValueError("interpolation can only be 'linear', 'lower', 'higher', "
                      "'midpoint', or 'nearest'")
   a, q = _promote_dtypes_inexact(a, q)
+  if issubdtype(a.dtype, np.complexfloating):
+    raise ValueError("quantile does not support complex input, as the operation is poorly defined.")
   if axis is None:
     a = ravel(a)
     axis = 0


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
@@ -1603,6 +1603,8 @@ def testPad(self, shape, dtype, mode, pad_width, constant_values):
       if (pad_width != () and stat_length != () and
           not (dtype in bool_dtypes and mode == 'mean'))))
   def testPadStatValues(self, shape, dtype, mode, pad_width, stat_length):
+    if mode == 'median' and np.issubdtype(dtype, np.complexfloating):
+      self.skipTest("median statistic is not supported for dtype=complex.")
     rng = jtu.rand_default(self.rng())
     args_maker = lambda: [rng(shape, dtype)]
Repository: google/jax

Issue:
Unexpected assertion error for gradient of ppermute.
The following code produces an assertion error:

```
import jax

def ppermute(input):
  return jax.lax.ppermute(input, axis_name="i", perm=[[0, 1], [1, 0]])

grad_fn = jax.grad(ppermute)

vmapped_gradients_fn = jax.vmap(grad_fn, axis_name="i")
  
vector = jax.numpy.array([1., 2.])
grads = vmapped_gradients_fn(vector)
```

```
AssertionError                            Traceback (most recent call last)
/usr/local/lib/python3.7/dist-packages/jax/_src/lax/parallel.py in _ppermute_batcher(axis_size, frame_name, _, vals_in, dims_in, axis_name, perm)
    791   assert axis_name[0] == frame_name, "ppermute batcher called with a wrong axis!"
    792   assert len(perm) == axis_size, "Permutation doesn't match the axis size!"
--> 793   assert d is not batching.not_mapped
    794   perm_indices = np.zeros(axis_size, dtype=int)
    795   for src, dst in perm:
```


Interestingly, the assertion error goes away, if I just inject a constant vector of ones:

```
import jax

def ppermute(input, one):
  return one * jax.lax.ppermute(input, axis_name="i", perm=[[0, 1], [1, 0]])

grad_fn = jax.grad(ppermute)

vmapped_gradients_fn = jax.vmap(grad_fn, axis_name="i")
  
vector = jax.numpy.array([1., 2.])
ones = jax.numpy.ones([2])
grads = vmapped_gradients_fn(vector, ones)
print(grads)  # [1. 1.]
```

This makes me believe the assert triggering above may indicate a bug.



Discussion:

Resolving patch:
diff --git a/jax/_src/lax/parallel.py b/jax/_src/lax/parallel.py
--- a/jax/_src/lax/parallel.py
+++ b/jax/_src/lax/parallel.py
@@ -793,7 +793,8 @@ def _ppermute_batcher(axis_size, frame_name, _, vals_in, dims_in, axis_name, per
     raise NotImplementedError("ppermute batcher only supports a single axis")
   assert axis_name[0] == frame_name, "ppermute batcher called with a wrong axis!"
   assert len(perm) == axis_size, "Permutation doesn't match the axis size!"
-  assert d is not batching.not_mapped
+  if d is batching.not_mapped:
+    return v, d
   perm_indices = np.zeros(axis_size, dtype=int)
   for src, dst in perm:
     perm_indices[dst] = src


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -1267,6 +1267,19 @@ def f(x):
     self.assertEqual(f(jnp.ones(3)).shape, (3,))
     self.assertEqual(jax.vmap(f)(jnp.ones((2, 3))).shape, (2, 3))
 
+  def testPpermuteBatcherTrivial(self):
+    # https://github.com/google/jax/issues/8688
+    def ppermute(input):
+      return jax.lax.ppermute(input, axis_name="i", perm=[[0, 1], [1, 0]])
+
+    grad_fn = jax.grad(ppermute)
+
+    vmapped_gradients_fn = jax.vmap(grad_fn, axis_name="i")
+
+    vector = jax.numpy.array([1., 2.])
+    ans = vmapped_gradients_fn(vector)  # doesn't crash
+    self.assertAllClose(ans, jnp.ones(2), check_dtypes=False)
+
 
 Array = Any
 ArrayElt = Any
Repository: google/jax

Issue:
jax.block_until_ready() for blocking until pytrees leaves are ready
It is sometimes useful (e.g., in benchmarking) to disable JAX's asynchronous dispatch: https://jax.readthedocs.io/en/latest/async_dispatch.html

The easy way to compute an array is to call `.block_until_ready()`, but this doesn't work on pytrees.

I propose adding a function `jax.block_until_ready()` for mapping block_until_ready() over the leaves of a pytree:
```python
def block_until_ready(pytree):
  return tree_util.tree_map(lambda x: x.block_until_ready(), pytree)
```
I've written versions of this helper function quite a few times and it would be nice not to have to do so again.

This version requires that all leaves are JAX arrays. Alternatively, we could skip calling `.block_until_ready()` on non-JAX arrays, which don't have the method defined. This might be more generally useful.


Discussion:
@mattjj or perhaps we should have a typeclass for things you can block on? 😄 

Resolving patch:
diff --git a/jax/__init__.py b/jax/__init__.py
--- a/jax/__init__.py
+++ b/jax/__init__.py
@@ -52,6 +52,7 @@
 )
 from jax._src.api import (
   ad,  # TODO(phawkins): update users to avoid this.
+  block_until_ready,
   checkpoint as checkpoint,
   checkpoint_policies as checkpoint_policies,
   closure_convert as closure_convert,
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -3138,3 +3138,22 @@ def invertible(fun: Callable) -> Callable:
     fun: The function assumed to be invertible.
   """
   return iad.invertible(fun)
+
+
+def block_until_ready(x):
+  """
+  Tries to call a ``block_until_ready`` method on pytree leaves.
+
+  Args:
+    x: a pytree, usually with at least some JAX array instances at its leaves.
+
+  Returns:
+    A pytree with the same structure and values of the input, where the values
+    of all JAX array leaves are ready.
+  """
+  def try_to_block(x):
+    try:
+      return x.block_until_ready()
+    except AttributeError:
+      return x
+  return jax.tree_util.tree_map(try_to_block, x)


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -1517,6 +1517,13 @@ def test_devicearray_block_until_ready(self):
     # Tests mostly that block_until_ready() does not produce an error.
     self.assertTrue(y is x)
 
+  def test_block_until_ready_function(self):
+    # Just tests that we don't error...
+    pytree = (device_put(1.), np.ones(3))
+    pytree = jax.block_until_ready(pytree)
+    self.assertAllClose(pytree[0], jnp.array(1.), check_dtypes=False)
+    self.assertAllClose(pytree[1], np.ones(3), check_dtypes=False)
+
   def test_devicearray_weakref_friendly(self):
     x = device_put(1.)
     y = weakref.ref(x)
Repository: google/jax

Issue:
set a max on the number of lines of tracer provenance printed
A user reported getting a ton of lines printed under "this value became a tracer due to JAX operations on these lines". That makes sense, since there's no limit on the amount of provenance information we attempt to print! Let's set a limit...


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
@@ -1159,11 +1159,13 @@ def _origin_msg(self):
       msts = ["  operation "
               f"{core.pp_eqn(eqn, core.JaxprPpContext(), print_shapes=True)}\n"
               f"    from line {source_info_util.summarize(eqn.source_info)}"
-              for eqn in progenitor_eqns]
+              for eqn in progenitor_eqns[:5]]  # show at most 5
       origin = (f"While tracing the function {dbg.func_src_info} "
                 f"for {dbg.traced_for}, "
                 "this value became a tracer due to JAX operations on these lines:"
                 "\n\n" + "\n\n".join(msts))
+      if len(progenitor_eqns) > 5:
+        origin += "\n\n(Additional originating lines are not shown.)"
     else:
       origin = (f"The error occured while tracing the function {dbg.func_src_info} "
                 f"for {dbg.traced_for}.")


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -2765,6 +2765,16 @@ def f():
     with self.assertRaisesRegex(core.ConcretizationTypeError, msg):
       f()
 
+  def test_concrete_error_because_const_2(self):
+    @jax.jit
+    def f():
+      result = sum(jnp.add(1, 1) for _ in range(6))
+      assert result > 0
+
+    msg = "Additional originating lines are not shown."
+    with self.assertRaisesRegex(core.ConcretizationTypeError, msg):
+      f()
+
   def test_xla_computation_zeros_doesnt_device_put(self):
     with jtu.count_device_put() as count:
       api.xla_computation(lambda: jnp.zeros(3))()
Repository: google/jax

Issue:
dynamic_update_slice start index adjustment fails when vmapped
The documentation for `dynamic_update_slice` explains "If the update slice is too large to fit in the array, the start index will be adjusted to make it fit" and this behavior does indeed seem to work normally but fails when I apply vmap. No update occurs if the update slice would not fit even if the update slice would still be partially in the array as in the example below.

```python
zeros = jnp.zeros(5)
ones = jnp.ones(3)

for i in range(1, 4):
    print(jax.lax.dynamic_update_slice(zeros, ones, (jnp.array(i, dtype=int),)))

print(jax.vmap(jax.lax.dynamic_update_slice, (None, None, 0))(zeros, ones, (jnp.arange(1, 4),)))

# [0. 1. 1. 1. 0.]
# [0. 0. 1. 1. 1.]
# [0. 0. 1. 1. 1.]
# [[0. 1. 1. 1. 0.]
#  [0. 0. 1. 1. 1.]
#  [0. 0. 0. 0. 0.]]
```


Discussion:
Thanks for the report; it looks like this is a bug in the `dynamic_update_slice` batching rule: https://github.com/google/jax/blob/04369a35888581b2c93abccffbe61fb3db8c460e/jax/_src/lax/slicing.py#L918-L936

On first look, I think the issue may be that it should use `mode=CLIP` rather than `mode=PROMISE_IN_BOUNDS`, but I'd have to dig deeper to be sure.

Resolving patch:
diff --git a/jax/_src/lax/slicing.py b/jax/_src/lax/slicing.py
--- a/jax/_src/lax/slicing.py
+++ b/jax/_src/lax/slicing.py
@@ -933,7 +933,7 @@ def _dynamic_update_slice_batching_rule(batched_args, batch_dims):
     scatter, (operand, index, update), (operand_bd, index_bdim, update_bd),
     update_jaxpr=None, update_consts=None, dimension_numbers=dnums,
     indices_are_sorted=True, unique_indices=True,
-    mode=GatherScatterMode.PROMISE_IN_BOUNDS)
+    mode=GatherScatterMode.CLIP)
 
 
 dynamic_update_slice_p = standard_primitive(


Test patch:
diff --git a/jax/experimental/jax2tf/tests/shape_poly_test.py b/jax/experimental/jax2tf/tests/shape_poly_test.py
--- a/jax/experimental/jax2tf/tests/shape_poly_test.py
+++ b/jax/experimental/jax2tf/tests/shape_poly_test.py
@@ -1701,6 +1701,9 @@ def _add_vmap_primitive_harnesses():
       "lu",
       "custom_linear_solve",
 
+      # Broken by https://github.com/google/jax/pull/9094
+      "dynamic_update_slice",
+
       # We do *= shapes in the batching rule for conv_general_dilated
       "conv_general_dilated",
 
diff --git a/tests/lax_test.py b/tests/lax_test.py
--- a/tests/lax_test.py
+++ b/tests/lax_test.py
@@ -1552,6 +1552,15 @@ def args_maker():
     self._CheckAgainstNumpy(lax_reference.dynamic_update_slice,
                             lax.dynamic_update_slice, args_maker)
 
+  def testDynamicUpdateSliceBatched(self):
+    # Regression test for https://github.com/google/jax/issues/9083
+    x = jnp.arange(5)
+    y = jnp.arange(6, 9)
+    ind = jnp.arange(6)
+    expected = jnp.vstack([lax.dynamic_update_slice(x, y, (i,)) for i in ind])
+    actual = jax.vmap(lax.dynamic_update_slice, (None, None, 0))(x, y, (ind,))
+    self.assertAllClose(expected, actual)
+
   @parameterized.named_parameters(jtu.cases_from_list(
       {"testcase_name": "_shape={}_perm={}".format(
           jtu.format_shape_dtype_string(shape, dtype), perm),
Repository: google/jax

Issue:
Some JAX flags cannot be set by environment variables
Currently, the `JAX_ENABLE_X64` flag can be set by environment variable:

```sh
$ export JAX_ENABLE_X64=True
$ python
>>> from jax.config import config
>>> config.FLAGS.jax_enable_x64
True
```

However, some flags cannot:

```sh
$ export JAX_XLA_BACKEND=tpu_driver
$ export JAX_BACKEND_TARGET="grpc://$COLAB_TPU_ADDR"
$ python
>>> from jax.config import config
>>> config.FLAGS.jax_xla_backend
''
>>> config.FLAGS.jax_backend_target
''
```


Discussion:
Thanks for raising this!

@skye since you were looking at this stuff recently, could you take a look at this issue?
@mattjj in which module should I look for this?
> in which module should I look for this?

@bharatnishant are you asking about these particular flags, or about configuration options in general?

In general, the 'good' config stuff is in jax/_src/config.py. That config stuff is good because it has a uniform way to set up environment variables, flags, and context managers. Some flags are defined elsewhere, and you can find those just by grepping for them.

The `jax_xla_backend` option is instead [defined in xla_bridge.py](https://github.com/google/jax/blob/68e9e1c26d5d9439d03d09a10b8f9b26e8258383/jax/_src/lib/xla_bridge.py#L57), but as you can see it's deprecated in favor of [`jax_platforms`, which does have an env var](https://github.com/google/jax/blob/68e9e1c26d5d9439d03d09a10b8f9b26e8258383/jax/_src/lib/xla_bridge.py#L69). @ayaka14732 I suggest using `export JAX_PLATFORMS=tpu_driver` instead of using `JAX_XLA_BACKEND`.

As for `JAX_BACKEND_TARGET`, it looks like that [doesn't have an environment variable](https://github.com/google/jax/blob/68e9e1c26d5d9439d03d09a10b8f9b26e8258383/jax/_src/lib/xla_bridge.py#L60). Maybe we should add one there...

I don't think #9118 touches `JAX_BACKEND_TARGET`, so I'll add an environment variable there now (in an ad-hoc way, rather than using general jax/_src/config.py stuff). @skye lmk if that's mistaken.

Resolving patch:
diff --git a/jax/_src/lib/xla_bridge.py b/jax/_src/lib/xla_bridge.py
--- a/jax/_src/lib/xla_bridge.py
+++ b/jax/_src/lib/xla_bridge.py
@@ -57,7 +57,8 @@
     'jax_xla_backend', '',
     'Deprecated, please use --jax_platforms instead.')
 flags.DEFINE_string(
-    'jax_backend_target', '',
+    'jax_backend_target',
+    os.getenv('JAX_BACKEND_TARGET', '').lower(),
     'Either "local" or "rpc:address" to connect to a remote service target.')
 # TODO(skye): warn when this is used once we test out --jax_platforms a bit
 flags.DEFINE_string(


Test patch: