Merge branch 'main' into agoose77/feat-expect-termination

src/awkward/operations/ak_nan_to_num.py

@@ -91,21 +91,17 @@ def action(layout, backend, **kwargs):
 
  def action(inputs, backend, **kwargs):
  if all(isinstance(x, ak.contents.NumpyArray) for x in inputs):
- tmp_layout = backend.nplike.asarray(inputs[0])
+ tmp_layout = inputs[0].data
  if id(nan) in broadcasting_ids:
- tmp_nan = backend.nplike.asarray(inputs[broadcasting_ids[id(nan)]])
+ tmp_nan = inputs[broadcasting_ids[id(nan)]].to_backend_array()
  else:
  tmp_nan = nan
  if id(posinf) in broadcasting_ids:
- tmp_posinf = backend.nplike.asarray(
- inputs[broadcasting_ids[id(posinf)]]
- )
+ tmp_posinf = inputs[broadcasting_ids[id(posinf)]].to_backend_array()
  else:
  tmp_posinf = posinf
  if id(neginf) in broadcasting_ids:
- tmp_neginf = backend.nplike.asarray(
- inputs[broadcasting_ids[id(neginf)]]
- )
+ tmp_neginf = inputs[broadcasting_ids[id(neginf)]].to_backend_array()
  else:
  tmp_neginf = neginf
  return (

tests/test_2591_nan_to_num.py

@@ -0,0 +1,15 @@
+# BSD 3-Clause License; see https://github.com/scikit-hep/awkward-1.0/blob/main/LICENSE
+
+import pytest
+
+import awkward as ak
+
+pytest.importorskip("jax")
+
+ak.jax.register_and_check()
+
+
+def test():
+ ak.nan_to_num(
+ ak.Array([1, 2, 3], backend="jax"), nan=ak.Array([1, 2, 3], backend="jax")
+ )