fix: assorted operation bugs (softmax, enforce_type, merge_union_of_records, and more)

src/awkward/_connect/rdataframe/from_rdataframe.py

@@ -63,7 +63,14 @@
 
 
 def from_rdataframe(
-    data_frame, columns, highlevel, behavior, with_name, offsets_type, keep_order
+    data_frame,
+    columns,
+    highlevel,
+    behavior,
+    attrs,
+    with_name,
+    offsets_type,
+    keep_order,
 ):
     if hasattr(data_frame, "proxied_node"):
         raise NotImplementedError("Distributed RDataFrame is not yet supported")
@@ -269,6 +276,7 @@ def form_dtype(form):
         depth_limit=1,
         highlevel=True,
         behavior=behavior,
+        attrs=attrs,
         with_name=with_name,
     )
 
@@ -278,6 +286,7 @@ def form_dtype(form):
             ak.contents.IndexedArray(sorted, out.layout),
             highlevel=True,
             behavior=behavior,
+            attrs=attrs,
         )
 
     return wrap_layout(out.layout, highlevel=highlevel, behavior=behavior)

src/awkward/operations/ak_argcartesian.py

@@ -118,7 +118,7 @@ def _impl(arrays, axis, nested, parameters, with_name, highlevel, behavior, attr
     if isinstance(arrays, Mapping):
         index_arrays = {n: ak.local_index(x, axis) for n, x in arrays.items()}
     else:
-        index_arrays = [ak.local_index(x) for x in arrays]
+        index_arrays = [ak.local_index(x, axis) for x in arrays]
 
     if with_name is not None:
         if parameters is None:

src/awkward/operations/ak_count_nonzero.py

@@ -98,8 +98,6 @@ def _impl(array, axis, keepdims, mask_identity, highlevel, behavior, attrs):
 
     axis = regularize_axis(axis, none_allowed=True)
 
-    with HighLevelContext(behavior=behavior, attrs=attrs) as ctx:
-        layout = ctx.unwrap(array, allow_record=False, primitive_policy="error")
     reducer = ak._reducers.CountNonzero()
 
     out = ak._do.reduce(

src/awkward/operations/ak_enforce_type.py

@@ -487,10 +487,8 @@ def _type_is_enforceable(
     elif layout.is_regular:
         if isinstance(type_, ak.types.RegularType):
             if layout.size == type_.size:
-                return _TypeEnforceableResult(
-                    is_enforceable=False, requires_packing=False
-                )
-            return _type_is_enforceable(layout.content, type_.content)
+                return _type_is_enforceable(layout.content, type_.content)
+            return _TypeEnforceableResult(is_enforceable=False, requires_packing=False)
 
         elif isinstance(type_, ak.types.ListType):
             return _type_is_enforceable(layout.content, type_.content)

src/awkward/operations/ak_from_raggedtensor.py

@@ -99,5 +99,5 @@ def _recursive_call(content, offsets_arr, count):
         )
     else:
         return ak.contents.ListOffsetArray(
-            offsets_arr[count], _recursive_call(content, offsets_arr, count)
+            offsets_arr[count], _recursive_call(content, offsets_arr, count + 1)
         )

src/awkward/operations/ak_from_rdataframe.py

@@ -53,11 +53,20 @@ def from_rdataframe(
 
     See also #ak.to_rdataframe.
     """
-    return _impl(rdf, columns, highlevel, behavior, with_name, offsets_type, keep_order)
+    return _impl(
+        rdf,
+        columns,
+        highlevel,
+        behavior,
+        attrs,
+        with_name,
+        offsets_type,
+        keep_order,
+    )
 
 
 def _impl(
-    data_frame, columns, highlevel, behavior, with_name, offsets_type, keep_order
+    data_frame, columns, highlevel, behavior, attrs, with_name, offsets_type, keep_order
 ):
     import awkward._connect.rdataframe.from_rdataframe  # noqa: F401
 
@@ -90,6 +99,7 @@ def _impl(
         columns,
         highlevel=highlevel,
         behavior=behavior,
+        attrs=attrs,
         with_name=with_name,
         offsets_type=offsets_type,
         keep_order=keep_order,

src/awkward/operations/ak_from_safetensors.py

@@ -109,7 +109,7 @@ def _impl(
         from safetensors import _safe_open_handle
     except ImportError as err:
         raise ImportError(
-            """to use ak.from_tensorflow, you must install the 'safetensors' package with:
+            """to use ak.from_safetensors, you must install the 'safetensors' package with:
 
         pip install safetensors
 or

src/awkward/operations/ak_merge_union_of_records.py

@@ -291,12 +291,12 @@ def apply(layout, depth, backend, **kwargs):
                     # Rewrite union index of indexed types
                     if content.is_indexed:
                         next_index_data[is_this_tag] = content.index.data[
-                            content.index.data[is_this_tag]
+                            layout.index.data[is_this_tag]
                         ]
                         next_contents.append(content.content)
 
                     else:
-                        next_index_data[is_this_tag] = content.index.data[is_this_tag]
+                        next_index_data[is_this_tag] = layout.index.data[is_this_tag]
                         next_contents.append(content)
 
                 return invert_record_union(

src/awkward/operations/ak_nan_to_num.py

@@ -50,7 +50,7 @@ def nan_to_num(
     See also #ak.nan_to_none to convert NaN to None, i.e. missing values with option-type.
     """
     # Dispatch
-    yield (array,)
+    yield array, nan, posinf, neginf
 
     # Implementation
     return _impl(array, copy, nan, posinf, neginf, highlevel, behavior, attrs)

src/awkward/operations/ak_round.py

@@ -40,8 +40,9 @@ def round(
         attrs (None or dict): Custom attributes for the output array, if
             high-level.
 
-    Returns the real components of the given array elements.
-    If the arrays have complex elements, the returned arrays are floats.
+    Rounds the given array elements to the given number of `decimals`.
+    Implements [np.round](https://numpy.org/doc/stable/reference/generated/numpy.round.html)
+    for Awkward Arrays.
     """
     # Dispatch
     yield (array,)

src/awkward/operations/ak_softmax.py

@@ -26,7 +26,7 @@
 @high_level_function()
 def softmax(
     x,
-    axis=None,
+    axis=-1,
     *,
     keepdims=False,
     mask_identity=False,
@@ -37,11 +37,12 @@ def softmax(
     """
     Args:
         x: The data on which to compute the softmax (anything #ak.to_layout recognizes).
-        axis (None or int or str): The dimension over which the softmax is
-            computed. If None, the softmax is computed over the innermost
-            dimension. If an int, `0` is the outermost dimension, `1` is the
-            first level of nested lists, etc., and negative values count from
-            the innermost: `-1` is the innermost, `-2` is the next level up, etc.
+        axis (int or str): The dimension over which the softmax is
+            computed (default is `-1`, the innermost dimension). ak.softmax is
+            currently only defined for the innermost axis. If an int, `0` is the
+            outermost dimension, `1` is the first level of nested lists, etc.,
+            and negative values count from the innermost: `-1` is the innermost,
+            `-2` is the next level up, etc.
             If a str, it is interpreted as the name of the axis which maps to
             an int if named axes are present.
             Named axes are attached to an array using #ak.with_named_axis and
@@ -96,6 +97,13 @@ def _impl(x, axis, keepdims, mask_identity, highlevel, behavior, attrs):
     axis = _named_axis_to_positional_axis(named_axis, axis)
     axis = regularize_axis(axis, none_allowed=True)
 
+    if axis is None:
+        raise ValueError(
+            "ak.softmax requires an explicit axis (it is only defined for "
+            "axis=-1); axis=None is not supported. "
+            "See https://github.com/scikit-hep/awkward/issues/2760"
+        )
+
     x = ctx.wrap(x_layout)
 
     if maybe_posaxis(x_layout, axis, 1) != maybe_posaxis(x_layout, -1, 1):

src/awkward/operations/ak_to_json.py

@@ -217,51 +217,64 @@ def opener():
             def opener():
                 return _NoContextManager(file)
 
-    try:
-        if line_delimited:
-            if file is None:
-                out = []
-                for datum in jsondata:
-                    out.append(
-                        json.dumps(
-                            datum,
-                            skipkeys=True,
-                            ensure_ascii=True,
-                            check_circular=False,
-                            allow_nan=False,
-                            indent=None,
-                            separators=separators,
-                            default=convert_other,
-                            sort_keys=False,
-                        )
+    if line_delimited:
+        if file is None:
+            out = []
+            for datum in jsondata:
+                out.append(
+                    json.dumps(
+                        datum,
+                        skipkeys=True,
+                        ensure_ascii=True,
+                        check_circular=False,
+                        allow_nan=False,
+                        indent=None,
+                        separators=separators,
+                        default=convert_other,
+                        sort_keys=False,
                     )
-                    out.append(line_delimited)
-                return "".join(out)
-
-            else:
-                with opener() as openfile:
-                    for datum in jsondata:
-                        json.dump(
-                            datum,
-                            openfile,
-                            skipkeys=True,
-                            ensure_ascii=True,
-                            check_circular=False,
-                            allow_nan=False,
-                            indent=None,
-                            separators=separators,
-                            default=convert_other,
-                            sort_keys=False,
-                        )
-                        openfile.write(line_delimited)
+                )
+                out.append(line_delimited)
+            return "".join(out)
 
         else:
-            if isinstance(array, (ak.highlevel.Record, ak.record.Record)):
-                jsondata = jsondata[0]
+            with opener() as openfile:
+                for datum in jsondata:
+                    json.dump(
+                        datum,
+                        openfile,
+                        skipkeys=True,
+                        ensure_ascii=True,
+                        check_circular=False,
+                        allow_nan=False,
+                        indent=None,
+                        separators=separators,
+                        default=convert_other,
+                        sort_keys=False,
+                    )
+                    openfile.write(line_delimited)
 
-            if file is None:
-                return json.dumps(
+    else:
+        if isinstance(array, (ak.highlevel.Record, ak.record.Record)):
+            jsondata = jsondata[0]
+
+        if file is None:
+            return json.dumps(
+                jsondata,
+                skipkeys=True,
+                ensure_ascii=True,
+                check_circular=False,
+                allow_nan=False,
+                indent=num_indent_spaces,
+                separators=separators,
+                default=convert_other,
+                sort_keys=False,
+            )
+        else:
+            with opener() as openfile:
+                return json.dump(
                     jsondata,
+                    openfile,
                     skipkeys=True,
                     ensure_ascii=True,
                     check_circular=False,
@@ -271,23 +284,6 @@ def opener():
                     default=convert_other,
                     sort_keys=False,
                 )
-            else:
-                with opener() as openfile:
-                    return json.dump(
-                        jsondata,
-                        openfile,
-                        skipkeys=True,
-                        ensure_ascii=True,
-                        check_circular=False,
-                        allow_nan=False,
-                        indent=num_indent_spaces,
-                        separators=separators,
-                        default=convert_other,
-                        sort_keys=False,
-                    )
-
-    except Exception as err:
-        raise err from err
 
 
 class _NoContextManager:

src/awkward/operations/ak_to_layout.py

@@ -259,7 +259,7 @@ def _impl(
             )
         else:
             raise ValueError(
-                f"Encountered an invalid string policy value {primitive_policy!r}. "
+                f"Encountered an invalid string policy value {string_policy!r}. "
                 f'The permitted values are "pass-through", "as-characters", "promote", and "error".'
             )
     elif isinstance(obj, (datetime, date, time, Number, bool)):

src/awkward/operations/ak_to_parquet_dataset.py

@@ -58,7 +58,7 @@ def _impl(directory, filenames, storage_options):
         ) from None
     fs, destination = fsspec.core.url_to_fs(directory, **(storage_options or {}))
     if not fs.isdir(destination):
-        raise ValueError(f"{destination!r} is not a directory" + {__file__})
+        raise ValueError(f"{destination!r} is not a directory")
 
     filepaths = get_filepaths(filenames, fs, destination)
 

src/awkward/operations/ak_unflatten.py

@@ -97,7 +97,7 @@ def unflatten(array, counts, axis=0, *, highlevel=True, behavior=None, attrs=Non
         See also #ak.num and #ak.flatten.
     """
     # Dispatch
-    yield (array,)
+    yield array, counts
 
     # Implementation
     return _impl(array, counts, axis, highlevel, behavior, attrs)

tests/test_2596_named_axis.py

@@ -691,19 +691,19 @@ def test_named_axis_ak_argcartesian():
 
     assert ak.argcartesian(
         [named_one, named_two, named_three], axis="x", nested=False
-    ).named_axis == {"x": 0, "y": 1}
+    ).named_axis == {"x": 0}
     assert ak.argcartesian(
         [named_one, named_two, named_three], axis="x", nested=True
-    ).named_axis == {"x": 1, "y": 2}
+    ).named_axis == {"x": 1}
     assert ak.argcartesian(
         [named_one, named_two, named_three], axis="x", nested=[0]
-    ).named_axis == {"x": 1, "y": 2}
+    ).named_axis == {"x": 1}
     assert ak.argcartesian(
         [named_one, named_two, named_three], axis="x", nested=[1]
-    ).named_axis == {"x": 0, "y": 2}
+    ).named_axis == {"x": 0}
     assert ak.argcartesian(
         [named_one, named_two, named_three], axis="x", nested=[0, 1]
-    ).named_axis == {"x": 2, "y": 3}
+    ).named_axis == {"x": 2}
 
 
 def test_negative_named_axis_ak_argcartesian():