broadcast ufuncs kwargs

Author: gdementen

larray/core/array.py

@@ -8541,6 +8541,18 @@ def raw_broadcastable(values, min_axes=None):
     return raw, res_axes
 
 
+def make_args_broadcastable(args, kwargs=None, min_axes=None):
+    """
+    Make args and kwargs (NumPy) broadcastable between them.
+    """
+    values = (args + tuple(kwargs.values())) if kwargs is not None else args
+    first_kw = len(args)
+    raw_bcast_values, res_axes = raw_broadcastable(values, min_axes=min_axes)
+    raw_bcast_args = raw_bcast_values[:first_kw]
+    raw_bcast_kwargs = dict(zip(kwargs.keys(), raw_bcast_values[first_kw:]))
+    return raw_bcast_args, raw_bcast_kwargs, res_axes
+
+
 _default_float_error_handler = float_error_handler_factory(3)
 
 

larray/core/ufuncs.py

@@ -3,15 +3,13 @@
 
 import numpy as np
 
-from larray.core.array import LArray, raw_broadcastable
+from larray.core.array import LArray, make_args_broadcastable
 
 
 def broadcastify(func):
     # intentionally not using functools.wraps, because it does not work for wrapping a function from another module
     def wrapper(*args, **kwargs):
-        # TODO: normalize args/kwargs like in LIAM2 so that we can also broadcast if args are given via kwargs
-        #       (eg out=)
-        raw_args, combined_axes = raw_broadcastable(args)
+        raw_bcast_args, raw_bcast_kwargs, res_axes = make_args_broadcastable(args, kwargs)
 
         # We pass only raw numpy arrays to the ufuncs even though numpy is normally meant to handle those cases itself
         # via __array_wrap__
@@ -25,9 +23,9 @@ def wrapper(*args, **kwargs):
         # It fails on "np.minimum(ndarray, LArray)" because it calls __array_wrap__(high, result) which cannot work if
         # there was broadcasting involved (high has potentially less labels than result).
         # it does this because numpy calls __array_wrap__ on the argument with the highest __array_priority__
-        res_data = func(*raw_args, **kwargs)
-        if combined_axes:
-            return LArray(res_data, combined_axes)
+        res_data = func(*raw_bcast_args, **raw_bcast_kwargs)
+        if res_axes:
+            return LArray(res_data, res_axes)
         else:
             return res_data
     # copy meaningful attributes (numpy ufuncs do not have __annotations__ nor __qualname__)