Merge pull request #20637 from jakevdp:array-api-scalar

PiperOrigin-RevId: 623184036

Author: jax authors

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):

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, /):