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

Issue:
Add factorial to `jax.scipy.special`
It would be useful to have the `scipy.special.factorial` function, even just to have the floating point approximate form from the gamma function.

```
def factorial(n):
    return gamma (n + 1)
```

Thanks!

Please:

- [x] Check for duplicate requests.
- [ ] Describe your goal, and if possible provide a code snippet with a motivating example.




Discussion:
Makes sense – we could probably implement `exact=True` as well by pre-computing the 13 or so factorials that are representable in int32.

Resolving patch:
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -58,6 +58,14 @@ def gamma(x: ArrayLike) -> Array:
 )(_betaln_impl)
 
 
+@_wraps(osp_special.factorial, module='scipy.special')
+def factorial(n: ArrayLike, exact: bool = False) -> Array:
+  if exact:
+    raise NotImplementedError("factorial with exact=True")
+  n, = promote_args_inexact("factorial", n)
+  return jnp.where(n < 0, 0, lax.exp(lax.lgamma(n + 1)))
+
+
 @_wraps(osp_special.beta, module='scipy.special')
 def beta(x: ArrayLike, y: ArrayLike) -> Array:
   x, y = promote_args_inexact("beta", x, y)
diff --git a/jax/scipy/special.py b/jax/scipy/special.py
--- a/jax/scipy/special.py
+++ b/jax/scipy/special.py
@@ -30,6 +30,7 @@
   expi as expi,
   expit as expit,
   expn as expn,
+  factorial as factorial,
   gammainc as gammainc,
   gammaincc as gammaincc,
   gammaln as gammaln,


Test patch:
diff --git a/tests/lax_scipy_special_functions_test.py b/tests/lax_scipy_special_functions_test.py
--- a/tests/lax_scipy_special_functions_test.py
+++ b/tests/lax_scipy_special_functions_test.py
@@ -89,6 +89,9 @@ def op_record(name, nargs, dtypes, rng_factory, test_grad, nondiff_argnums=(), t
     op_record(
         "gammaln", 1, float_dtypes, jtu.rand_positive, False
     ),
+    op_record(
+        "factorial", 1, float_dtypes, jtu.rand_default, True
+    ),
     op_record(
         "i0", 1, float_dtypes, jtu.rand_default, True
     ),
Repository: google/jax

Issue:
`debug.print` looks strange in jaxpr representation
```python
import jax
def f(x):
  jax.debug.print("{}", x)
  return x
jax.make_jaxpr(f)(1)
```
```
{ lambda ; a:i32[]. let
     = debug_callback[
      callback=<function debug_callback.<locals>._flat_callback at 0x7c30cf833ac0>
      effect=Debug
    ] a
  in (a,) }
```
Should there be a variable (or a `_`) in front of the `=` sign? Or should the `=` sign be removed?


Discussion:

Resolving patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -3086,7 +3086,10 @@ def _pp_eqn(eqn, context, settings) -> pp.Doc:
   rhs = [pp.text(eqn.primitive.name, annotation=name_stack_annotation),
          pp_kv_pairs(sorted(eqn.params.items()), context, settings),
          pp.text(" ") + pp_vars(eqn.invars, context)]
-  return pp.concat([lhs, pp.text(" = ", annotation=annotation), *rhs])
+  if lhs.format():
+    return pp.concat([lhs, pp.text(" = ", annotation=annotation), *rhs])
+  else:
+    return pp.concat(rhs)
 CustomPpEqnRule = Callable[[JaxprEqn, JaxprPpContext, JaxprPpSettings], pp.Doc]
 pp_eqn_rules: dict[Primitive, CustomPpEqnRule]  = {}
 


Test patch:
diff --git a/tests/api_test.py b/tests/api_test.py
--- a/tests/api_test.py
+++ b/tests/api_test.py
@@ -6197,6 +6197,14 @@ def test_convert_element_type_literal_constant_folding(self):
     jaxpr = api.make_jaxpr(lambda: cet(3.))()
     self.assertLen(jaxpr.eqns, 0)
 
+  def test_eqn_repr_with_no_lhs(self):
+    def f(x):
+      jax.debug.print("{}", x)
+      return x
+    jaxpr = jax.make_jaxpr(f)(np.int32(0))
+    self.assertEqual(jaxpr.eqns[0].primitive, jax._src.debugging.debug_callback_p)
+    self.assertStartsWith(str(jaxpr.eqns[0]), "debug_callback[", )
+
 
 class DCETest(jtu.JaxTestCase):
Repository: google/jax

Issue:
Shard map + CumSum/conv_general_dilated missing replication rule
I stumbled upon 

> NotImplementedError: No replication rule for cumsum. As a workaround, pass the `check_rep=False` argument to `shard_map`. To get this fixed, open an issue at https://github.com/google/jax/issues (edited) 

And in another code

> No replication rule for conv_general_dilated. As a workaround, pass the `check_rep=False` argument to `shard_map`. To get this fixed, open an issue at


Do you need reproducers?






Discussion:
Nope, I suspect the error is self-explanatory to @mattjj 
Thanks for reporting this! No need for a reproducer. I'm surprised I missed conv_general_dilated...
What's your JAX version? I think #18740 added `conv_general_dilated` two weeks ago on HEAD.

Resolving patch:
diff --git a/jax/experimental/shard_map.py b/jax/experimental/shard_map.py
--- a/jax/experimental/shard_map.py
+++ b/jax/experimental/shard_map.py
@@ -917,6 +917,12 @@ def _standard_collective_rewrite(prim, mesh, in_rep, x, axis_name, **params):
     register_standard_check(o)
     register_standard_rewrite(o)
 
+for p in [control_flow.loops.cumsum_p, control_flow.loops.cumlogsumexp_p,
+          control_flow.loops.cumprod_p, control_flow.loops.cummax_p,
+          control_flow.loops.cummin_p]:
+  register_standard_check(p)
+  register_standard_rewrite(p)
+
 
 @register_check(lax_parallel.psum_p)
 def _psum_check(_, *in_rep, axes, axis_index_groups):


Test patch:
diff --git a/tests/shard_map_test.py b/tests/shard_map_test.py
--- a/tests/shard_map_test.py
+++ b/tests/shard_map_test.py
@@ -1294,6 +1294,12 @@ def f(x, y):
     y = f(a, b)  # don't crash
     self.assertAllClose(y, a @ b, check_dtypes=False, atol=1e-2, rtol=1e-2)
 
+  def test_cumsum(self):
+    mesh = jtu.create_global_mesh((4,), ('i',))
+    x = jnp.arange(8.)
+    shard_map(jnp.cumsum, mesh=mesh, in_specs=P('i'), out_specs=P('i')
+              )(x)  # don't crash
+
   def test_custom_jvp_inside_jit(self):
     mesh = jtu.create_global_mesh((4,), ('batch',))
     x = shard_map(jax.jit(jax.nn.relu),
Repository: google/jax

Issue:
Unexpected behavior of `jax.scipy.stats.binom.pmf`
### Description

pmf of a random variable should be zero outside of its range. While plotting the graph for `jax.scipy.stats.binom.pmf`, I notice that for $n>5$ and $p>0.5$, there are some oscillations in the values of the pmf, which should not be there. For evidence, I am attaching a plot too.

```python
import jax
from jax import numpy as jnp
from matplotlib import pyplot as plt

x = jnp.linspace(-1, 10, 1000)
xxf = jax.scipy.stats.binom.pmf(k=x, n=5, p=0.8)

plt.plot(x, xxf)
plt.tight_layout()
plt.show()
```
![xyxy](https://github.com/google/jax/assets/62182585/c35a2ebb-e914-4b37-9f59-6ccd7cc10b28)
The side left to the zero is as expected.

### What jax/jaxlib version are you using?

jax v0.4.23

### Which accelerator(s) are you using?

CPU


Discussion:

Resolving patch:
diff --git a/jax/_src/scipy/stats/binom.py b/jax/_src/scipy/stats/binom.py
--- a/jax/_src/scipy/stats/binom.py
+++ b/jax/_src/scipy/stats/binom.py
@@ -33,7 +33,7 @@ def logpmf(k: ArrayLike, n: ArrayLike, p: ArrayLike, loc: ArrayLike = 0) -> Arra
     )
     log_linear_term = lax.add(xlogy(y, p), xlog1py(lax.sub(n, y), lax.neg(p)))
     log_probs = lax.add(comb_term, log_linear_term)
-    return jnp.where(lax.lt(k, loc), -jnp.inf, log_probs)
+    return jnp.where(lax.ge(k, loc) & lax.lt(k, loc + n + 1), log_probs, -jnp.inf)
 
 
 @_wraps(osp_stats.nbinom.pmf, update_doc=False)


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
@@ -1159,6 +1159,10 @@ def args_maker():
                               tol=5e-4)
       self._CompileAndCheck(lax_fun, args_maker, rtol=tol, atol=tol)
 
+  def testBinomPmfOutOfRange(self):
+    # Regression test for https://github.com/google/jax/issues/19150
+    self.assertEqual(lsp_stats.binom.pmf(k=6.5, n=5, p=0.8), 0.0)
+
   def testIssue972(self):
     self.assertAllClose(
       np.ones((4,), np.float32),
Repository: google/jax

Issue:
Concatenation error in `jax.numpy.diff`
### Description

Consider the following example:

```python
import jax.numpy as jnp

x = jnp.arange(10)
z = jnp.diff(x, prepend=x[0])
```

The snippet above will throw an error: 

```python
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/.../miniconda3/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 694, in diff
    arr = concatenate(combined, axis)
  File "/home/.../miniconda3/lib/python3.10/site-packages/jax/_src/numpy/lax_numpy.py", line 1841, in concatenate
    raise ValueError("Zero-dimensional arrays cannot be concatenated.")
ValueError: Zero-dimensional arrays cannot be concatenated.
```

Such an error does not appear for `numpy`. That is, `z = np.diff(x, prepend=x[0])` will work in numpy. Hence, I believe that this is a design bug in jax.numpy. The behaviou does not match the API description.

### What jax/jaxlib version are you using?

0.4.23

### Which accelerator(s) are you using?

CPU and GPU

### Additional system info?

Ubuntu

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report! This seems like a case we missed in the implementation.

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
@@ -670,7 +670,7 @@ def diff(a: ArrayLike, n: int = 1, axis: int = -1,
   combined: list[Array] = []
   if prepend is not None:
     util.check_arraylike("diff", prepend)
-    if isscalar(prepend):
+    if not ndim(prepend):
       shape = list(arr.shape)
       shape[axis] = 1
       prepend = broadcast_to(prepend, tuple(shape))
@@ -680,7 +680,7 @@ def diff(a: ArrayLike, n: int = 1, axis: int = -1,
 
   if append is not None:
     util.check_arraylike("diff", append)
-    if isscalar(append):
+    if not ndim(append):
       shape = list(arr.shape)
       shape[axis] = 1
       append = broadcast_to(append, tuple(shape))


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
@@ -2744,6 +2744,16 @@ def np_fun(x, n=n, axis=axis, prepend=prepend, append=append):
     self._CheckAgainstNumpy(np_fun, jnp_fun, args_maker, check_dtypes=False)
     self._CompileAndCheck(jnp_fun, args_maker)
 
+  def testDiffPrepoendScalar(self):
+    # Regression test for https://github.com/google/jax/issues/19362
+    x = jnp.arange(10)
+    result_jax = jnp.diff(x, prepend=x[0], append=x[-1])
+
+    x = np.array(x)
+    result_numpy = np.diff(x, prepend=x[0], append=x[-1])
+
+    self.assertArraysEqual(result_jax, result_numpy)
+
   @jtu.sample_product(
     op=["zeros", "ones"],
     shape=[2, (), (2,), (3, 0), np.array((4, 5, 6), dtype=np.int32),
Repository: google/jax

Issue:
jax.lax.rng_uniform(k,k) has different behaviour in CPU (RANDOM int32) /TPU (ALWAYS k)
### Description

rng_uniform claims "Returns uniformly distributed random numbers in the range [a, b)".

When calling it with a == b, it has a different behaviour from CPU and TPU:
  - CPU: it seems to use the full int32 range.
  - TPU: it ALWAYS returns the same!!!

Example:

```
import jax.tools.colab_tpu
# jax.tools.colab_tpu.setup_tpu()

import jax
from jax.lib import xla_bridge
print(xla_bridge.get_backend().platform)
print(jax.devices())
print(jax.lax.rng_uniform(3, 3, (10,)))
```

CPU:
```
cpu
[CpuDevice(id=0)]
[-1752798562  1744106001  -688905896  -376700610    28523420  -600052118
  1113020891   874027959   762190079  -467109515]
```

GPU:
```
gpu
[cuda(id=0)]
[-1752798562  1744106001  -688905896  -376700610    28523420  -600052118
  1113020891   874027959   762190079  -467109515]
```

On TPU:

```
tpu
[TpuDevice(id=0, process_index=0, coords=(0,0,0), core_on_chip=0), TpuDevice(id=1, process_index=0, coords=(0,0,0), core_on_chip=1), TpuDevice(id=2, process_index=0, coords=(1,0,0), core_on_chip=0), TpuDevice(id=3, process_index=0, coords=(1,0,0), core_on_chip=1), TpuDevice(id=4, process_index=0, coords=(0,1,0), core_on_chip=0), TpuDevice(id=5, process_index=0, coords=(0,1,0), core_on_chip=1), TpuDevice(id=6, process_index=0, coords=(1,1,0), core_on_chip=0), TpuDevice(id=7, process_index=0, coords=(1,1,0), core_on_chip=1)]
[3 3 3 3 3 3 3 3 3 3]
```

### What jax/jaxlib version are you using?

0.4.23

### Which accelerator(s) are you using?

CPU/GPU/TPU

### Additional system info?

https://colab.research.google.com/

### NVIDIA GPU info

_No response_


Discussion:
Thanks for the report! I think this is working as expected. From the [XLA:RngUniform docs](https://www.tensorflow.org/xla/operation_semantics#rnguniform):

> if *b <= a* the result is implementation-defined.

I think we could make that more clear in the JAX documentation for `jax.lax.rng_uniform`.

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
@@ -4437,7 +4437,9 @@ def _outfeed_lowering(ctx, token, *xs, partitions):
 def rng_uniform(a, b, shape):
   """Stateful PRNG generator. Experimental and its use is discouraged.
 
-  Returns uniformly distributed random numbers in the range [a, b)
+  Returns uniformly distributed random numbers in the range [a, b). If
+  b <= a, then the result is undefined, and different implementations may
+  return different results.
 
   You should use jax.random for most purposes; this function exists only for
   niche use cases with special performance requirements.


Test patch:
Repository: google/jax

Issue:
Documentation Errors in jax.scipy module
I've found the following errors in the documentation for [`jax.scipy`](https://jax.readthedocs.io/en/latest/jax.scipy.html) module:
https://jax.readthedocs.io/en/latest/jax.scipy.html

1. Duplicate `sqrtm` function:

- [ ] The `sqrtm` function is listed twice under `jax.scipy.linalg`.

2. Typo in idctn description:

- [ ] The description for `idctn` incorrectly states "Discrete Cosine Transform." It should be "Inverse Discrete Cosine Transform" to match the function's behavior.




Discussion:

Resolving patch:
diff --git a/jax/_src/scipy/fft.py b/jax/_src/scipy/fft.py
--- a/jax/_src/scipy/fft.py
+++ b/jax/_src/scipy/fft.py
@@ -109,7 +109,7 @@ def dctn(x: Array, type: int = 2,
   return x
 
 
-@implements(osp_fft.dct)
+@implements(osp_fft.idct)
 def idct(x: Array, type: int = 2, n: int | None = None,
         axis: int = -1, norm: str | None = None) -> Array:
   if type != 2:


Test patch:
Repository: google/jax

Issue:
Not captured error in `jnp.broadcast_to`
### Description

If `shape` is shorter than `array.ndim`, then an exception in `safe_zip` is not captured, which is a bit confusing.

Example:

```python
import jax.numpy as jnp
jnp.broadcast_to(jnp.zeros((3, 4, 5)), (4, 5))
```

JAX gives
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "jax/_src/numpy/lax_numpy.py", line 1222, in broadcast_to
    return util._broadcast_to(array, shape)
  File "jax/_src/numpy/util.py", line 425, in _broadcast_to
    for arr_d, shape_d in safe_zip(arr_shape, shape_tail))
ValueError: safe_zip() argument 2 is shorter than argument 1
```

While numpy gives the message
```
ValueError: input operand has more dimensions than allowed by the axis remapping
```

### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.25
jaxlib: 0.4.25
numpy:  1.24.3
python: 3.10.13 (main, Sep 11 2023, 13:44:35) [GCC 11.2.0]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
platform: uname_result(system='Linux', node='<HIDDEN>', release='5.15.0-78-generic', version='#85-Ubuntu SMP Fri Jul 7 15:25:09 UTC 2023', machine='x86_64')
```


Discussion:

Resolving patch:
diff --git a/jax/_src/numpy/util.py b/jax/_src/numpy/util.py
--- a/jax/_src/numpy/util.py
+++ b/jax/_src/numpy/util.py
@@ -418,6 +418,8 @@ def _broadcast_to(arr: ArrayLike, shape: DimSize | Shape) -> Array:
   arr_shape = np.shape(arr)
   if core.definitely_equal_shape(arr_shape, shape):
     return arr
+  elif len(shape) < len(arr_shape):
+    raise ValueError(f"Cannot broadcast to shape with fewer dimensions: {arr_shape=} {shape=}")
   else:
     nlead = len(shape) - len(arr_shape)
     shape_tail = shape[nlead:]


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
@@ -5177,6 +5177,13 @@ def testBroadcastTo(self, from_shape, to_shape):
     self._CheckAgainstNumpy(np_op, jnp_op, args_maker)
     self._CompileAndCheck(jnp_op, args_maker)
 
+  def testBroadcastToInvalidShape(self):
+    # Regression test for https://github.com/google/jax/issues/20533
+    x = jnp.zeros((3, 4, 5))
+    with self.assertRaisesRegex(
+        ValueError, "Cannot broadcast to shape with fewer dimensions"):
+      jnp.broadcast_to(x, (4, 5))
+
   @jtu.sample_product(
     [dict(shapes=shapes, broadcasted_shape=broadcasted_shape)
       for shapes, broadcasted_shape in [
Repository: google/jax

Issue:
jax.scipy.special.beta results are different from scipy.special.beta for certain input
### Description

When I use` jax.scipy.special.beta`, I get different results from the scipy one for some input as below:

```
b = 14.25
n = b - 1
k = jnp.arange(40)
plt.plot(k, scipy.special.beta(1 + n - k, 1 + k), label="scipy", ls=":")
plt.plot(k, jax.scipy.special.beta(1 + n - k, 1 + k), label="jax")
plt.plot(k, np.array([cephes4py.beta(1 + n - k, 1 + k) for k in np.arange(40)]), label="cephes", ls="-.")
plt.legend();
plt.show();
```
Here is the result:

![image](https://github.com/google/jax/assets/8914509/c0cfd3c2-1027-4d2e-b59e-412d90d44237)


### System info (python version, jaxlib version, accelerator, etc.)

```
jax:    0.4.26
jaxlib: 0.4.25
numpy:  1.26.4
python: 3.10.13
```


Discussion:
Thanks for the report! It looks like we weren't properly setting the sign of `beta` for negative inputs. #20589 should fix this.

Resolving patch:
diff --git a/jax/_src/scipy/special.py b/jax/_src/scipy/special.py
--- a/jax/_src/scipy/special.py
+++ b/jax/_src/scipy/special.py
@@ -44,14 +44,16 @@ def gammaln(x: ArrayLike) -> Array:
   return lax.lgamma(x)
 
 
+def _gamma_sign(x: Array) -> Array:
+  floor_x = lax.floor(x)
+  return jnp.where((x > 0) | (x == floor_x) | (floor_x % 2 == 0), 1.0, -1.0)
+
+
 @implements(osp_special.gamma, module='scipy.special', lax_description="""\
 The JAX version only accepts real-valued inputs.""")
 def gamma(x: ArrayLike) -> Array:
   x, = promote_args_inexact("gamma", x)
-  # Compute the sign for negative x, matching the semantics of scipy.special.gamma
-  floor_x = lax.floor(x)
-  sign = jnp.where((x > 0) | (x == floor_x), 1.0, (-1.0) ** floor_x)
-  return sign * lax.exp(lax.lgamma(x))
+  return _gamma_sign(x) * lax.exp(lax.lgamma(x))
 
 betaln = implements(
     osp_special.betaln,
@@ -71,7 +73,8 @@ def factorial(n: ArrayLike, exact: bool = False) -> Array:
 @implements(osp_special.beta, module='scipy.special')
 def beta(x: ArrayLike, y: ArrayLike) -> Array:
   x, y = promote_args_inexact("beta", x, y)
-  return lax.exp(betaln(x, y))
+  sign = _gamma_sign(x) * _gamma_sign(y) * _gamma_sign(x + y)
+  return sign * lax.exp(betaln(x, y))
 
 
 @implements(osp_special.betainc, module='scipy.special')


Test patch:
diff --git a/tests/lax_scipy_special_functions_test.py b/tests/lax_scipy_special_functions_test.py
--- a/tests/lax_scipy_special_functions_test.py
+++ b/tests/lax_scipy_special_functions_test.py
@@ -53,10 +53,10 @@ def op_record(name, nargs, dtypes, rng_factory, test_grad, nondiff_argnums=(), t
 
 JAX_SPECIAL_FUNCTION_RECORDS = [
     op_record(
-        "beta", 2, float_dtypes, jtu.rand_positive, False
+        "beta", 2, float_dtypes, jtu.rand_default, False
     ),
     op_record(
-        "betaln", 2, float_dtypes, jtu.rand_positive, False
+        "betaln", 2, float_dtypes, jtu.rand_default, False
     ),
     op_record(
         "betainc", 3, float_dtypes, jtu.rand_positive, False
Repository: google/jax

Issue:
Tracer's imag method returns float; crashes with the Array API
### Description
The following:
```python
import jax.numpy as jnp
from jax.lax import while_loop
from array_api_compat import get_namespace

def body(x):
    assert False

def cond(x):
    # x is a tracer, x.real is a tracer, x.imag is a float!
    xp = get_namespace(x)
    xp.square(x.imag)  # No problem with x.real!

while_loop(cond, body, jnp.asarray(0.0))
```
crashes with:
```
ValueError: square: inputs must be arrays; got types [<class 'float'>]
```

Also broken:
```python
import jax.numpy as jnp
from jax import jit

@jit
def f(x):
    x.imag.shape  # Crash!

f(jnp.asarray(0.0))
```
According to the Array API spec, `x.imag` must be a _Jax array_ (not a NumPy array, nor a float).

### System info (python version, jaxlib version, accelerator, etc.)

An NVIDIA GPU may be present on this machine, but a CUDA-enabled jaxlib is not installed. Falling back to cpu.
jax:    0.4.26
jaxlib: 0.4.26
numpy:  1.26.4
python: 3.11.8 (main, Feb 22 2024, 17:25:49) [GCC 11.4.0]


Discussion:
Thanks - it seems that the validation for array API inputs is too strict: we should probably allow Python scalars here: https://github.com/google/jax/blob/1d26365e601528d4cb63ca9d39213248eb320adb/jax/experimental/array_api/_elementwise_functions.py#L25-L26

Resolving patch:
diff --git a/jax/experimental/array_api/_data_type_functions.py b/jax/experimental/array_api/_data_type_functions.py
--- a/jax/experimental/array_api/_data_type_functions.py
+++ b/jax/experimental/array_api/_data_type_functions.py
@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-
+import builtins
 import functools
 from typing import NamedTuple
 import jax
@@ -183,6 +183,8 @@ def isdtype(dtype, kind):
 def result_type(*arrays_and_dtypes):
   dtypes = []
   for val in arrays_and_dtypes:
+    if isinstance(val, (builtins.bool, int, float, complex)):
+      val = jax.numpy.array(val)
     if isinstance(val, jax.Array):
       val = val.dtype
     if _is_valid_dtype(val):
diff --git a/jax/experimental/array_api/_elementwise_functions.py b/jax/experimental/array_api/_elementwise_functions.py
--- a/jax/experimental/array_api/_elementwise_functions.py
+++ b/jax/experimental/array_api/_elementwise_functions.py
@@ -22,10 +22,11 @@
 
 def _promote_dtypes(name, *args):
   assert isinstance(name, str)
-  if not all(isinstance(arg, jax.Array) for arg in args):
+  if not all(isinstance(arg, (bool, int, float, complex, jax.Array))
+             for arg in args):
     raise ValueError(f"{name}: inputs must be arrays; got types {[type(arg) for arg in args]}")
   dtype = _result_type(*args)
-  return [arg.astype(dtype) for arg in args]
+  return [jax.numpy.asarray(arg).astype(dtype) for arg in args]
 
 
 def abs(x, /):


Test patch: