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

Issue:
eval_shape appears to hang when tracing a large random sample
Calling

```python
import jax
from jax import random
jax.eval_shape(lambda: random.normal(key=random.PRNGKey(0), shape=[int(1e11)]))
```

returns `ShapeDtypeStruct(shape=(100000000000,), dtype=float32)`. I'd expect that calling

```python
jax.eval_shape(lambda: random.normal(key=random.PRNGKey(0), shape=[int(1e11), int(1e11)]))
```

would similarly return `ShapeDtypeStruct(shape=(100000000000,100000000000), dtype=float32)` (of course, if we *actually* ran the function we'd get an OOM error). 

But instead, this call seems to end up in an ~infinite loop (at least, it hangs for > 2 minutes). Reproduction here:
https://colab.research.google.com/drive/14IZF3WUTaq976jaUC5tehKTPwt3YhNTe?usp=sharing

Is this expected?



Discussion:
It looks like this comes from the implementation of `_random_bits`, specifically this part: https://github.com/google/jax/blob/1d0c461ea67db683206d2643a5b377f2f8ff2e84/jax/_src/random.py#L306-L312

If the number of generated numbers is larger than the largest uint32, the implementation stacks multiple calls to the underlying random number generator in a Python loop. Since these calls are required for abstract evaluation, they are generated even for `jax.eval_shape`. And if the size is large enough, the abstract evaluation graph will become quite large and slow to evaluate.

You can see this if you wish by calling `make_jaxpr()` for your input.

We could maybe think about fixing this (perhaps compute it via `vmap`?), but it does seem like a pretty rare corner case. In what context did you encounter the problem?

Resolving patch:
diff --git a/jax/_src/random.py b/jax/_src/random.py
--- a/jax/_src/random.py
+++ b/jax/_src/random.py
@@ -304,14 +304,15 @@ def _random_bits(key, bit_width, shape):
   max_count = int(np.ceil(bit_width * size / 32))
 
   nblocks, rem = divmod(max_count, jnp.iinfo(np.uint32).max)
+
   if not nblocks:
     bits = threefry_2x32(key, lax.iota(np.uint32, rem))
   else:
-    *subkeys, last_key = split(key, nblocks + 1)
-    blocks = [threefry_2x32(k, lax.iota(np.uint32, jnp.iinfo(np.uint32).max))
-              for k in subkeys]
+    keys = split(key, nblocks + 1)
+    subkeys, last_key = keys[:-1], keys[-1]
+    blocks = vmap(threefry_2x32, in_axes=(0, None))(subkeys, lax.iota(np.uint32, jnp.iinfo(np.uint32).max))
     last = threefry_2x32(last_key, lax.iota(np.uint32, rem))
-    bits = lax.concatenate(blocks + [last], 0)
+    bits = lax.concatenate([blocks.ravel(), last], 0)
 
   dtype = _UINT_DTYPES[bit_width]
   if bit_width == 64:


Test patch:
Repository: google/jax

Issue:
jax.numpy.transpose does not sufficiently validate inputs
For example:
```python
>>> import numpy as np
>>> import jax.numpy as jnp
>>> x = np.zeros((2, 3, 4))

# jnp.transpose is silent on invalid permutations
>>> jnp.transpose(x, (0, 1, 2, 3))
array([[[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]],

       [[0., 0., 0., 0.],
        [0., 0., 0., 0.],
        [0., 0., 0., 0.]]])

# numpy transpose raises an error.
>>> np.transpose(x, (0, 1, 2, 3))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<__array_function__ internals>", line 5, in transpose
  File "/Users/vanderplas/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 658, in transpose
    return _wrapfunc(a, 'transpose', axes)
  File "/Users/vanderplas/.local/share/virtualenvs/jax-LBbfM5ix/lib/python3.8/site-packages/numpy/core/fromnumeric.py", line 58, in _wrapfunc
    return bound(*args, **kwds)
ValueError: axes don't match array
```


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
@@ -1216,7 +1216,7 @@ def transpose(operand: Array, permutation: Sequence[int]) -> Array:
   operator.
   """
   permutation = tuple(permutation)
-  if permutation == tuple(range(len(permutation))):
+  if permutation == tuple(range(np.ndim(operand))):
     return operand
   else:
     return transpose_p.bind(operand, permutation=permutation)


Test patch:
Repository: google/jax

Issue:
Constraints for jaxlib version during jax installation
In v0.2.11 jax added the nice optional installation targets `cpu`, `cudaXXX` and `tpu` which install jaxlib for the desired platform. Currently, the setup script guarantees that the installed version of jaxlib is compatible by constraining the version number from below to the earliest known compatible version. However, no similar constraint is put for a maximum compatible version number. Given that jaxlib does not appear to necessarily maintain backward compatibility between updates, this is an issue for installing older versions of jax (as these will always pull the latest available jaxlib version, leading to crashes due to incompatibility).

Jax installation should always result in a working setup, i.e., only the latest compatible version of jaxlib should be installed when installing jax. This either requires version constraints in jax's setup.py or backward compatibility guarantees in jaxlib (or a combination of both).

Please:

- [x] Check for duplicate requests.
- [x] Describe your goal, and if possible provide a code snippet with a motivating example.



Discussion:
Thanks for the report – I believe we do pin jaxlib to the current version at the time of the given JAX release for GPU and TPU installation; see the source here: https://github.com/google/jax/blob/c97d63dec33845a13c9243a8f821a7d093ff5cfa/setup.py#L43-L60

That said, pip's behavior surrounding extra requirements and version specifications has surprised me before, so it may be that pip doesn't respect these version pinnings during actual installation, but I would consider that a pip bug.
A potential TODO is to pin the extra requirements for cpu to `jaxlib==_current_jaxlib_version` rather than `jaxlib>=_minimum_jaxlib_version`. Would that fix your issue?
Ah, I was apparently focusing only on the `cpu` target and didn't notice that for GPU and TPU this is indeed pinned. The todo you suggest would prevent the incompatibility problem I was worrying about, at the (small?) cost of also preventing users from benefitting from updates to jaxlib that would still be compatible to the version of jax they are installing. But I think this latter point is something that cannot be helped without some backward-compatibility policy reflected in version numbering for jaxlib/jax, so for now I'd settle for this pinning.
Sounds good - are you interested in sending a PR? If not I'm happy to take care of it.
I don't have the time right now, so please go ahead :)

Resolving patch:
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -46,7 +46,7 @@
 
         # CPU-only jaxlib can be installed via:
         # $ pip install jax[cpu]
-        'cpu': [f'jaxlib>={_minimum_jaxlib_version}'],
+        'cpu': [f'jaxlib=={_current_jaxlib_version}'],
 
         # Cloud TPU VM jaxlib can be installed via:
         # $ pip install jax[tpu] -f https://storage.googleapis.com/jax-releases/jax_releases.html


Test patch:
Repository: google/jax

Issue:
Document jax.device_get
The function `jax.device_get` is not documented. Preferably, it should appear on this page
https://jax.readthedocs.io/en/latest/jax.html

The code for the function also does not provide any doc-string.
https://github.com/google/jax/blob/26e9ebcdae5e610d1173fac1bc3504a9c3f2011f/jax/_src/api.py#L2258


Discussion:

Resolving patch:
diff --git a/jax/_src/api.py b/jax/_src/api.py
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2328,7 +2328,35 @@ def _device_get(x):
   else:
     return copy()
 
-def device_get(x):
+def device_get(x: Any):
+  """Transfer ``x`` to host.
+
+  Args:
+    x: An array, scalar, DeviceArray or (nested) standard Python container thereof
+      representing the array to be transferred to host.
+
+  Returns:
+    An array or (nested) Python container thereof representing the
+    value of ``x``.
+
+  Examples:
+    Passing a DeviceArray:
+
+    >>> import jax
+    >>> x = jax.numpy.array([1., 2., 3.])
+    >>> jax.device_get(x)
+    array([1., 2., 3.], dtype=float32)
+
+    Passing a scalar (has no effect):
+
+    >>> jax.device_get(1)
+    1
+
+  See Also:
+    - device_put
+    - device_put_sharded
+    - device_put_replicated
+  """
   for y in tree_leaves(x):
     try:
       y.copy_to_host_async()


Test patch:
Repository: google/jax

Issue:
jax.scipy.stats.gamma.pdf(x, a=1.) handles x=0 incorrectly
```python
scipy.stats.gamma.pdf(0., 1.)  # = exp(0.)
# 1.0
jax.scipy.stats.gamma.pdf(0., 1.)
# DeviceArray(nan, dtype=float32)
```


Discussion:
It's probably a simple matter of not using `xlogy` in the implementation, I can look later.

Resolving patch:
diff --git a/jax/_src/scipy/stats/gamma.py b/jax/_src/scipy/stats/gamma.py
--- a/jax/_src/scipy/stats/gamma.py
+++ b/jax/_src/scipy/stats/gamma.py
@@ -18,7 +18,7 @@
 from jax._src.numpy.util import _wraps
 from jax._src.numpy.lax_numpy import (_promote_args_inexact, _constant_like,
                                       where, inf)
-from jax.scipy.special import gammaln
+from jax.scipy.special import gammaln, xlogy
 
 
 @_wraps(osp_stats.gamma.logpdf, update_doc=False)
@@ -26,7 +26,7 @@ def logpdf(x, a, loc=0, scale=1):
   x, a, loc, scale = _promote_args_inexact("gamma.logpdf", x, a, loc, scale)
   one = _constant_like(x, 1)
   y = lax.div(lax.sub(x, loc), scale)
-  log_linear_term = lax.sub(lax.mul(lax.sub(a, one), lax.log(y)), y)
+  log_linear_term = lax.sub(xlogy(lax.sub(a, one), y), y)
   shape_terms = lax.add(gammaln(a), lax.log(scale))
   log_probs = lax.sub(log_linear_term, shape_terms)
   return where(lax.lt(x, loc), -inf, log_probs)


Test patch:
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -231,6 +231,11 @@ def args_maker():
                             tol=5e-4)
     self._CompileAndCheck(lax_fun, args_maker)
 
+  def testGammaLogPdfZero(self):
+    # Regression test for https://github.com/google/jax/issues/7256
+    self.assertAllClose(
+      osp_stats.gamma.pdf(0.0, 1.0), lsp_stats.gamma.pdf(0.0, 1.0), atol=1E-6)
+
   @genNamedParametersNArgs(3)
   def testLaplaceLogPdf(self, shapes, dtypes):
     rng = jtu.rand_positive(self.rng())
Repository: google/jax

Issue:
ravel, flatten don't return copies of 1D arrays
The docstring for ravel states:
```
Return a contiguous flattened array.

LAX-backend implementation of :func:`ravel`.

The JAX version of this function will return a copy rather than a view of the input.
```

But this isn't true if the array is already flattened:
```python
import jax.numpy as jnp
x = jnp.array([1, 2, 3])
x.ravel().unsafe_buffer_pointer() == x.unsafe_buffer_pointer() #False
```
Why does this matter? It probably doesn't, but the behavior doesn't match the documentation.


Discussion:

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
@@ -1103,7 +1103,8 @@ def _sinc_maclaurin_jvp(k, primals, tangents):
   return _sinc_maclaurin(k, x), _sinc_maclaurin(k + 1, x) * t
 
 _ARRAY_VIEW_DOC = """
-The JAX version of this function will return a copy rather than a view of the input.
+The JAX version of this function may in some cases return a copy rather than a
+view of the input.
 """
 
 @_wraps(np.transpose, lax_description=_ARRAY_VIEW_DOC)


Test patch:
Repository: google/jax

Issue:
num replicas not correctly detected with `lax.switch` / `lax.cond`
Thanks to @LenaMartens for the repro!

```python
import os
os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=32'

import jax

def identity(x):
  return x

def cond_of_pmap(x):
  y = jax.lax.cond(True, lambda x: jax.pmap(identity)(x), lambda x: jax.pmap(identity)(x), x)
  return y

cond_of_pmap(jax.numpy.zeros((8, 2)))
```

```
AssertionError                            Traceback (most recent call last)
google3/third_party/py/jax/interpreters/xla.py in axis_groups(axis_env, name)
    533   mesh_axes = tuple(unsafe_map(partial(axis_read, axis_env), name))
    534   trailing_size, ragged = divmod(axis_env.nreps, prod(axis_env.sizes))
--> 535   assert not ragged
    536   mesh_spec = axis_env.sizes + (trailing_size,)
    537   return _axis_groups(mesh_spec, mesh_axes)

AssertionError: 
```

The issue was in `core.traverse_jaxpr_params`, which assumed that all subjaxprs appear as values in the params dict, while `lax.switch` (aka `lax.cond`) actually has a param value which is a _tuple_ of jaxprs. Those tuple elements were being ignored!

I'm going to make a quick fix which just flattens params to one level, but I think a better approach here would be to have a primitive-keyed table of rules for traversing subjaxprs.


Discussion:

Resolving patch:
diff --git a/jax/core.py b/jax/core.py
--- a/jax/core.py
+++ b/jax/core.py
@@ -307,11 +307,13 @@ def extract_call_jaxpr(
     return (params["call_jaxpr"], new_params)
 
 
+# TODO(mattjj): replace this approach with a primitive-keyed table of rules
 def traverse_jaxpr_params(f, params):
   """Applies f to each jaxpr parameter and returns a tuple of returned values."""
-  return {name: f(param)
+  return {name: f(p)
           for name, param in params.items()
-          if type(param) in (Jaxpr, ClosedJaxpr)}
+          for p in (param if isinstance(param, (tuple, list)) else [param])
+          if type(p) in (Jaxpr, ClosedJaxpr)}
 
 
 def eval_jaxpr(jaxpr: Jaxpr, consts, *args):


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -1701,6 +1701,17 @@ def func(_):
     out_dtype = func(unused_arg).dtype
     self.assertEqual(out_dtype, dtype)
 
+  def test_num_replicas_with_switch(self):
+    # https://github.com/google/jax/issues/7411
+    def identity(x):
+      return x
+
+    def cond_of_pmap(x):
+      y = lax.cond(True, jax.pmap(identity), jax.pmap(identity), x)
+      return y
+
+    cond_of_pmap(jnp.zeros((xla_bridge.device_count(), 2)))
+
 
 class VmapOfPmapTest(jtu.JaxTestCase):
Repository: google/jax

Issue:
Use bazel --version to check the bazel version
Please use `bazel --version` to check the version of Bazel. The currently used variant ignores the bazel-startup-options and actually starts a full bazel server which is a) overkill and b) wrong in case of options like `--output-user-root` and no or readonly home folders

Potential change:
```
try:
  version_output = shell([bazel_path, "--version"])
except subprocess.CalledProcessError:
  return False
match = re.search("(Build label:|bazel) *([0-9\\.]+)", version_output)
```

I kept the regexp generic enough so that you can still fall back to the current `bazel version` call and parse both kinds of output


Discussion:

Resolving patch:
diff --git a/build/build.py b/build/build.py
--- a/build/build.py
+++ b/build/build.py
@@ -209,7 +209,7 @@ def get_bazel_path(bazel_path_flag):
 
 def check_bazel_version(bazel_path):
   try:
-    version_output = shell([bazel_path, "--bazelrc=/dev/null", "version"])
+    version_output = shell([bazel_path, "--version"])
   except subprocess.CalledProcessError:
     return False
   match = re.search("Build label: *([0-9\\.]+)[^0-9\\.]", version_output)


Test patch:
Repository: google/jax

Issue:
script "examples/advi.py" throws "ValueError" exception
Run:
 ```python3 jax/examples/advi.py```
Output:
```
Optimizing variational parameters...
Iteration 0 lower bound 0.4957694113254547
Traceback (most recent call last):
  File "jax/examples/advi.py", line 138, in <module>
    callback(params, t)
  File "jax/examples/advi.py", line 98, in callback
    X, Y, Z = mesh_eval(target_dist, x_limits, y_limits, 1)
  File "jax/examples/advi.py", line 67, in mesh_eval
    return _mesh_eval(func, x_limits, y_limits, params, num_ticks)
ValueError: Non-hashable static arguments are not supported. An error occured while trying to hash an object of type <class 'list'>, [-2, 2]. The error was:
TypeError: unhashable type: 'list'
```





Discussion:

Resolving patch:
diff --git a/examples/advi.py b/examples/advi.py
--- a/examples/advi.py
+++ b/examples/advi.py
@@ -86,8 +86,8 @@ def objective(params, t):
     ax = fig.add_subplot(111, frameon=False)
     plt.ion()
     plt.show(block=False)
-    x_limits = [-2, 2]
-    y_limits = [-4, 2]
+    x_limits = (-2, 2)
+    y_limits = (-4, 2)
     target_dist = lambda x, _: jnp.exp(funnel_log_density(x))
     approx_dist = lambda x, params: jnp.exp(diag_gaussian_logpdf(x, *params))
 


Test patch:
Repository: google/jax

Issue:
jax.scipy.stats.beta.pdf handles x=0, x=1 incorrectly
```
a = b = 1.
x = np.array([0., 1.])
scipy.stats.beta.pdf(x, a, b)
# array([1., 1.])
jax.scipy.stats.beta.pdf(x, a, b)
# DeviceArray([nan, nan], dtype=float32) 
```

Same issue and resolution as #7256 basically--use xlogy/xlog1py here:
https://github.com/google/jax/blob/f6f1debf700b55ad8c0ecc1daa0ed1f677c2bf00/jax/_src/scipy/stats/beta.py#L30-L31




Discussion:
Thanks for the report!

Resolving patch:
diff --git a/jax/_src/scipy/stats/beta.py b/jax/_src/scipy/stats/beta.py
--- a/jax/_src/scipy/stats/beta.py
+++ b/jax/_src/scipy/stats/beta.py
@@ -18,7 +18,7 @@
 from jax._src.numpy.util import _wraps
 from jax._src.numpy.lax_numpy import (_promote_args_inexact, _constant_like,
                                       where, inf, logical_or)
-from jax.scipy.special import betaln
+from jax.scipy.special import betaln, xlogy, xlog1py
 
 
 @_wraps(osp_stats.beta.logpdf, update_doc=False)
@@ -27,8 +27,8 @@ def logpdf(x, a, b, loc=0, scale=1):
   one = _constant_like(x, 1)
   shape_term = lax.neg(betaln(a, b))
   y = lax.div(lax.sub(x, loc), scale)
-  log_linear_term = lax.add(lax.mul(lax.sub(a, one), lax.log(y)),
-                            lax.mul(lax.sub(b, one), lax.log1p(lax.neg(y))))
+  log_linear_term = lax.add(xlogy(lax.sub(a, one), y),
+                            xlog1py(lax.sub(b, one), lax.neg(y)))
   log_probs = lax.sub(lax.add(shape_term, log_linear_term), lax.log(scale))
   return where(logical_or(lax.gt(x, lax.add(loc, scale)),
                           lax.lt(x, loc)), -inf, log_probs)


Test patch:
diff --git a/tests/scipy_stats_test.py b/tests/scipy_stats_test.py
--- a/tests/scipy_stats_test.py
+++ b/tests/scipy_stats_test.py
@@ -148,6 +148,13 @@ def args_maker():
     self._CompileAndCheck(lax_fun, args_maker,
                           rtol={np.float32: 2e-3, np.float64: 1e-4})
 
+  def testBetaLogPdfZero(self):
+    # Regression test for https://github.com/google/jax/issues/7645
+    a = b = 1.
+    x = np.array([0., 1.])
+    self.assertAllClose(
+      osp_stats.beta.pdf(x, a, b), lsp_stats.beta.pdf(x, a, b), atol=1E-6)
+
   @genNamedParametersNArgs(3)
   def testCauchyLogPdf(self, shapes, dtypes):
     rng = jtu.rand_default(self.rng())
Repository: google/jax

Issue:
Jax PRNGKeyArray.shape is a method, not a property
Was this an intended change? 

```python
>>> import jax
>>> jax.config.update('jax_enable_custom_prng', True)
>>> k=jax.random.PRNGKey(0)
>>> k.shape
<bound method PRNGKeyArray.shape of PRNGKeyArray:
  shape = ()
  impl = PRNGImpl:
           key_shape = (2,)
           seed = <function threefry_seed at 0x12e6f00d0>
           split = <function threefry_split at 0x10cc2aaf0>
           random_bits = <CompiledFunction object at 0x12e6c8880>
           fold_in = <function threefry_fold_in at 0x10cc2a790>>
>>> k.shape()
()
```


Discussion:

Resolving patch:
diff --git a/jax/_src/prng.py b/jax/_src/prng.py
--- a/jax/_src/prng.py
+++ b/jax/_src/prng.py
@@ -125,6 +125,7 @@ def dtype(self):
           'deprecated `dtype` attribute of PRNG key arrays', FutureWarning)
       return np.uint32
 
+  @property
   def shape(self):
     # TODO(frostig): simplify once we always enable_custom_prng
     if config.jax_enable_custom_prng:


Test patch:
diff --git a/tests/random_test.py b/tests/random_test.py
--- a/tests/random_test.py
+++ b/tests/random_test.py
@@ -1095,6 +1095,22 @@ class LaxRandomWithCustomPRNGTest(LaxRandomTest):
   def seed_prng(self, seed):
     return prng.seed_with_impl(double_threefry_prng_impl, seed)
 
+  def test_split_shape(self):
+    key = self.seed_prng(73)
+    keys = random.split(key, 10)
+    self.assertEqual(keys.shape, (10,))
+
+  def test_vmap_fold_in_shape(self):
+    key = self.seed_prng(73)
+    keys = vmap(lambda i: random.fold_in(key, i))(jnp.arange(3))
+    self.assertEqual(keys.shape, (3,))
+
+  def test_cannot_add(self):
+    key = self.seed_prng(73)
+    self.assertRaisesRegex(
+        TypeError, r'unsupported operand type\(s\) for \+*',
+        lambda: key + 47)
+
 def _sampler_unimplemented_with_custom_prng(*args, **kwargs):
   raise SkipTest('sampler only implemented for default RNG')
Repository: google/jax

Issue:
Cuda Configuration Error: Invalid compute capability: compute_70build:rocm
> python3 build/build.py --enable_cuda

```
Cuda Configuration Error: Invalid compute capability: compute_70build:rocm
INFO: Found applicable config definition build:cuda in file /home/skoonce/jax/jax/.bazelrc: --repo_env TF_NEED_CUDA=1 --action_env TF_CUDA_COMPUTE_CAPABILITIES=3.5,5.2,6.0,6.1,7.0 --crosstool_top=@local_config_cuda//crosstool:toolchain --@local_config_cuda//:enable_cuda --define=xla_python_enable_gpu=true
ERROR: @local_config_cuda//:enable_cuda :: Error loading option @local_config_cuda//:enable_cuda: 
Cuda Configuration Error: Invalid compute capability: compute_70build:rocm

b''
Traceback (most recent call last):
  File "build/build.py", line 521, in <module>
    main()
  File "build/build.py", line 516, in main
    shell(command)
  File "build/build.py", line 53, in shell
    output = subprocess.check_output(cmd)
  File "/usr/lib/python3.8/subprocess.py", line 415, in check_output
    return run(*popenargs, stdout=PIPE, timeout=timeout, check=True,
  File "/usr/lib/python3.8/subprocess.py", line 516, in run
    raise CalledProcessError(retcode, process.args,
subprocess.CalledProcessError: Command '['./bazel-4.1.0-linux-x86_64', 'run', '--verbose_failures=true', '--config=avx_posix', '--config=mkl_open_source_only', '--config=cuda', ':build_wheel', '--', '--output_path=/home/skoonce/jax/jax/dist', '--cpu=x86_64']' returned non-zero exit status 2.
```

ping @yashk2810 eg d0acd9f3435a7eff924c9be937dc2abb1716e79a


Discussion:
Ahh, there's a `\n` missing. I'll fix it when I have access to my computer or you can send a PR too 😃
@yashk2810 kk, give me a sec! 🤓
Can you test it too with the command you ran? 

Thank you!

Resolving patch:
diff --git a/build/build.py b/build/build.py
--- a/build/build.py
+++ b/build/build.py
@@ -248,13 +248,13 @@ def write_bazelrc(python_bin_path=None, remote_build=None,
               .format(cudnn_version=cudnn_version))
     if cuda_compute_capabilities:
       f.write(
-        f'build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="{cuda_compute_capabilities}"')
+        f'build:cuda --action_env TF_CUDA_COMPUTE_CAPABILITIES="{cuda_compute_capabilities}"\n')
     if rocm_toolkit_path:
       f.write("build --action_env ROCM_PATH=\"{rocm_toolkit_path}\"\n"
               .format(rocm_toolkit_path=rocm_toolkit_path))
     if rocm_amdgpu_targets:
       f.write(
-        f'build:rocm --action_env TF_ROCM_AMDGPU_TARGETS="{rocm_amdgpu_targets}"')
+        f'build:rocm --action_env TF_ROCM_AMDGPU_TARGETS="{rocm_amdgpu_targets}"\n')
     if cpu is not None:
       f.write("build --distinct_host_configuration=true\n")
       f.write(f"build --cpu={cpu}\n")


Test patch: