Skip to content

Refactors 28 functions into numpy's ufunc style#5423

Draft
drculhane wants to merge 3 commits intoBears-R-Us:mainfrom
drculhane:moreUfunc
Draft

Refactors 28 functions into numpy's ufunc style#5423
drculhane wants to merge 3 commits intoBears-R-Us:mainfrom
drculhane:moreUfunc

Conversation

@drculhane
Copy link
Contributor

@drculhane drculhane commented Feb 19, 2026

This may never get reviewed or merged, but it was important to me to complete it.

28 functions (abs, fabs, ceil, floor, trunc, round, sign, isfinite, isinf, isnan, sin, cos, tan, arcsin, arccos, arctan, sinh, cosh, tanh, arcsinh, arccosh, arctanh, log, log2, log10, log1p, exp, expm1) have been rewritten into a framework that supports the out and where parameters of numpy ufuncs.

The new code starts in the general vicinity of line 2979 in numpy/numeric.py, and continues for the rest of the file. Although it's a lot of text, there's also a lot of structure to it.

This class:

class UfuncSpec:
    name: str  # Python-facing name
    chapel_name: str  # used in cmd=...
    argname: str  # "x" or "pda" in args dict
    output_dtype_resolver: Callable[[str, Any], Any]
    scalar_op: Optional[Callable[..., Any]] = None
    validate: Optional[ValidateHook] = None
    precompute: Optional[PrecomputeHook] = None
    extra_args_builder: Optional[ExtraArgsBuilder] = None
    input_cast: Optional[InputCastHook] = None
    chapel_accepts_dtype: bool = True  # isfinite/isinf/isnan float-only path uses False

is used to tell the API:

  • the function's name
  • how it's known chapel-side
  • whether its argument is called "x" or "pda" chapel-side (in a future life, I'll standardize this)
  • how to resolve the output dtype
  • what numpy scalar function to use if the input is scalar and out is None
  • what function (validate) to use to resolve special case dypes (e.g. some treat bool as int, some as float)
  • what precompute function to use if there's no need to call the server (e.g. uint supplied to abs)
  • how to handle non-standards args (this is how round's "decimal" is handled)
  • a special case flag for the isfinite/isinf/isnan functions where the new interface doesn't specify dtype

Here's an example of a validation and a precompute:

def validate_sign(name: str, ops: Sequence[Any]) -> None:
    (x,) = ops
    if isinstance(x, pdarray):
        if x.dtype == ak_bool:
            raise TypeError("sign does not support bool input")
    else:
        if resolve_scalar_dtype(x) == "bool":
            raise TypeError("sign does not support bool input")
    

def abs_precompute(ops: Sequence[pdarray], res_dtype: Any) -> Optional[pdarray]:
    (bx,) = ops
    if bx.dtype in (ak_bool, ak_uint64):
        return bx.copy()
    return None

A UfuncSpec object is created for each of the 28 functions.

Each of the 28 functions calls ufunc_unary something like this:

return ufunc_unary(ABS_SPEC, x, out=out, where=where)

Then ufunc_unary does all of the dtype handling, broadcasting of dimensions as needed, casting of types as needed, and implementation of where and out.

It passes all the unit tests, but this has been a big enough project that I may very well have missed things.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant