import inspect

import re

from functools import lru_cache, reduce

from operator import mul

from types import FunctionType

import numpy as np

from ... import _STANDARD_OPERATOR_NAMES, backend, binary, monoid, op

from ...dtypes import (

BOOL,

FP32,

FP64,

INT8,

INT16,

INT32,

INT64,

UINT8,

UINT16,

UINT32,

UINT64,

_supports_complex,

lookup_dtype,

)

from ...exceptions import UdfParseError, check_status_carg

from .. import _has_numba, _supports_udfs, ffi, lib

from ..dtypes import _sample_values

from ..expr import InfixExprBase

from .base import (

_SS_OPERATORS,

OpBase,

ParameterizedUdf,

TypedOpBase,

_call_op,

_hasop,

)

# Imported unconditionally (plain dict, no numba): ``_compile_udt`` consults it

# even when numba is absent.

from .udt_utils import BUILTIN_UDT_BINARY_OPS as _BUILTIN_UDT_BINARY_OPS

if _has_numba:

import numba

from .base import (

_bool_to_int8,

_compile_udf_for_udt,

_finalize_udt_op,

_get_udt_wrapper,

_resolve_udt_return_type,

)

if _supports_complex:

from ...dtypes import FC32, FC64

ffi_new = ffi.new

if _has_numba:

def _make_udt_comparison(dtype, dtype2, *, is_eq):

"""Build a cfunc-ready wrapper for UDT eq or ne comparison.

Compares per-leaf with IEEE semantics (NaN compares unequal even

to itself); scalar broadcasts to every leaf when paired with a

UDT. Mismatched-shape UDT pairs raise ``KeyError``. Returns

``(wrapper_func, wrapper_sig)``.

"""

from .udt_utils import _check_udt_pair, _get_udt_info, _iter_record_leaves

info_x = _get_udt_info(dtype)

info_y = _get_udt_info(dtype2)

op_name = "eq" if is_eq else "ne"

_check_udt_pair(op_name, dtype, dtype2, info_x, info_y)

x_is_scalar = info_x is None

y_is_scalar = info_y is None

udt_dtype = dtype2 if x_is_scalar else dtype

udt_info = info_y if x_is_scalar else info_x

nt = numba.types

cmp_op = "==" if is_eq else "!="

join = " and " if is_eq else " or "

if udt_info[0] == "array":

# Array UDT: unroll element-by-element so the scalar broadcast

# case can drop ``numba.carray`` on that side entirely.

# ``_get_udt_info`` already flattens the multi-dim shape.

# The wrapper sees the UDT side as a flat pointer-to-element,

# not pointer-to-Record; passing the Record numba_type here

# makes ``numba.carray`` try to read N record-sized chunks.

_base, N = udt_info[1]

base_numba = numba.from_dtype(_base)

x_ptr_type = (

nt.CPointer(numba.from_dtype(dtype.np_type))

if x_is_scalar

else nt.CPointer(base_numba)

)

y_ptr_type = (

nt.CPointer(numba.from_dtype(dtype2.np_type))

if y_is_scalar

else nt.CPointer(base_numba)

)

wrapper_sig = nt.void(nt.CPointer(INT8.numba_type), x_ptr_type, y_ptr_type)

if x_is_scalar:

terms = [f"(x_ptr[0] {cmp_op} y[{i}])" for i in range(N)]

arrays = f" y = numba.carray(y_ptr, {N})\n"

elif y_is_scalar:

terms = [f"(x[{i}] {cmp_op} y_ptr[0])" for i in range(N)]

arrays = f" x = numba.carray(x_ptr, {N})\n"

else:

terms = [f"(x[{i}] {cmp_op} y[{i}])" for i in range(N)]

arrays = f" x = numba.carray(x_ptr, {N})\n y = numba.carray(y_ptr, {N})\n"

body = join.join(terms)

src = f"def wrapper(z_ptr, x_ptr, y_ptr):\n{arrays} z_ptr[0] = {body}\n"

wrapper = _compile_codegen(

src,

func_name="wrapper",

source_label=f"<gb-udt {op_name} array N={N}>",

)

return wrapper, wrapper_sig

# Record UDT (possibly nested): chain leaf comparisons. For an

# array-valued sub-field, unroll its elements so the scalar side

# never sees a numpy fancy-broadcast (Numba's record access can't

# express that as a single expression).

np_type = udt_dtype.np_type

terms = []

for py_path, _c, leaf_dtype in _iter_record_leaves(np_type):

if leaf_dtype.subdtype is not None:

_base, shape = leaf_dtype.subdtype

sub_N = int(reduce(mul, shape))

if x_is_scalar:

subterms = [f"(x_ptr[0] {cmp_op} y[0]{py_path}[{i}])" for i in range(sub_N)]

elif y_is_scalar:

subterms = [f"(x[0]{py_path}[{i}] {cmp_op} y_ptr[0])" for i in range(sub_N)]

else:

reducer = ".all()" if is_eq else ".any()"

terms.append(f"(x[0]{py_path} {cmp_op} y[0]{py_path}){reducer}")

continue

terms.append("(" + join.join(subterms) + ")")

else:

x_expr = "x_ptr[0]" if x_is_scalar else f"x[0]{py_path}"

y_expr = "y_ptr[0]" if y_is_scalar else f"y[0]{py_path}"

terms.append(f"({x_expr} {cmp_op} {y_expr})")

expr = join.join(terms) if terms else ("True" if is_eq else "False")

lines = ["def wrapper(z_ptr, x_ptr, y_ptr):"]

if not x_is_scalar:

lines.append(" x = numba.carray(x_ptr, 1)")

if not y_is_scalar:

lines.append(" y = numba.carray(y_ptr, 1)")

lines.append(f" z_ptr[0] = {expr}")

src = "\n".join(lines) + "\n"

wrapper = _compile_codegen(

src,

func_name="wrapper",

source_label=f"<gb-udt {op_name} record nleaves={len(terms)}>",

)

x_ptr_type = (

nt.CPointer(numba.from_dtype(dtype.np_type))

if x_is_scalar

else nt.CPointer(dtype.numba_type)

)

y_ptr_type = (

nt.CPointer(numba.from_dtype(dtype2.np_type))

if y_is_scalar

else nt.CPointer(dtype2.numba_type)

)

wrapper_sig = nt.void(nt.CPointer(INT8.numba_type), x_ptr_type, y_ptr_type)

return wrapper, wrapper_sig

class TypedBuiltinBinaryOp(TypedOpBase):

__slots__ = ()

opclass = "BinaryOp"

def __call__(self, left, right=None, *, left_default=None, right_default=None):

if left_default is not None or right_default is not None:

if (

left_default is None

or right_default is None

or right is not None

or not isinstance(left, InfixExprBase)

or left.method_name != "ewise_add"

):

raise TypeError(

"Specifying `left_default` or `right_default` keyword arguments implies "

"performing `ewise_union` operation with infix notation.\n"

"There is only one valid way to do this:\n\n"

f">>> {self}(x | y, left_default=0, right_default=0)\n\nwhere x and y "

"are Vectors or Matrices, and left_default and right_default are scalars."

)

return left.left._ewise_union(

left.right, self, left_default, right_default, is_infix=True

)

return _call_op(self, left, right)

@property

def monoid(self):

rv = getattr(monoid, self.name, None)

if rv is not None and self.type in rv._typed_ops:

return rv[self.type]

@property

def commutes_to(self):

commutes_to = self.parent.commutes_to

if commutes_to is not None and (self.type in commutes_to._typed_ops or self.type._is_udt):

return commutes_to[self.type]

@property

def _semiring_commutes_to(self):

commutes_to = self.parent._semiring_commutes_to

if commutes_to is not None and (self.type in commutes_to._typed_ops or self.type._is_udt):

return commutes_to[self.type]

@property

def is_commutative(self):

return self.commutes_to is self

@property

def type2(self):

return self.type if self._type2 is None else self._type2

class TypedUserBinaryOp(TypedOpBase):

__slots__ = "_monoid"

opclass = "BinaryOp"

_owns_gb_obj = True

def __init__(self, parent, name, type_, return_type, gb_obj, dtype2=None):

super().__init__(parent, name, type_, return_type, gb_obj, f"{name}_{type_}", dtype2=dtype2)

self._monoid = None

@property

def monoid(self):

if self._monoid is None:

monoid = self.parent.monoid

if monoid is not None and self.type in monoid:

self._monoid = monoid[self.type]

return self._monoid

@property

def orig_func(self):

return self.parent.orig_func

@property

def _numba_func(self):

return self.parent._numba_func

commutes_to = TypedBuiltinBinaryOp.commutes_to

_semiring_commutes_to = TypedBuiltinBinaryOp._semiring_commutes_to

is_commutative = TypedBuiltinBinaryOp.is_commutative

type2 = TypedBuiltinBinaryOp.type2

__call__ = TypedBuiltinBinaryOp.__call__

class ParameterizedBinaryOp(ParameterizedUdf):

__slots__ = "func", "__signature__", "_monoid", "_cached_call", "_commutes_to", "_is_udt"

def __init__(self, name, func, *, anonymous=False, is_udt=False):

self.func = func

self.__signature__ = inspect.signature(func)

self._monoid = None

self._is_udt = is_udt

if name is None:

name = getattr(func, "__name__", name)

super().__init__(name, anonymous)

method = self._call_to_cache.__get__(self, type(self))

self._cached_call = lru_cache(maxsize=1024)(method)

self.__call__ = self._call

self._commutes_to = None

def _call_to_cache(self, *args, **kwargs):

binary = self.func(*args, **kwargs)

binary._parameterized_info = (self, args, kwargs)

return BinaryOp.register_anonymous(binary, self.name, is_udt=self._is_udt)

def _call(self, *args, **kwargs):

binop = self._cached_call(*args, **kwargs)

if self._monoid is not None and binop._monoid is None:

# This is all a bit funky. We try our best to associate a binaryop

# to a monoid. So, if we made a ParameterizedMonoid using this object,

# then try to create a monoid with the given arguments.

binop._monoid = binop # temporary!

try:

# If this call is successful, then it will set `binop._monoid`

self._monoid(*args, **kwargs) # pylint: disable=not-callable

except Exception:

binop._monoid = None

# assert binop._monoid is not binop

if self.is_commutative:

binop._commutes_to = binop

# Don't bother yet with creating `binop.commutes_to` (but we could!)

return binop

@property

def monoid(self):

return self._monoid

@property

def commutes_to(self):

if isinstance(self._commutes_to, str):

self._commutes_to = BinaryOp._find(self._commutes_to)

return self._commutes_to

is_commutative = TypedBuiltinBinaryOp.is_commutative

def _floordiv(x, y):

return x // y # pragma: no cover (numba)

def _rfloordiv(x, y):

return y // x # pragma: no cover (numba)

def _absfirst(x, y):

return np.abs(x) # pragma: no cover (numba)

def _abssecond(x, y):

return np.abs(y) # pragma: no cover (numba)

def _rpow(x, y):

return y**x # pragma: no cover (numba)

def _isclose(rel_tol=1e-7, abs_tol=0.0):

def inner(x, y): # pragma: no cover (numba)

return x == y or abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol)

return inner

_MAX_INT64 = np.iinfo(np.int64).max

def _binom(N, k): # pragma: no cover (numba)

# Returns 0 if overflow or out-of-bounds

if k > N or k < 0:

return 0

val = np.int64(1)

for i in range(min(k, N - k)):

if val > _MAX_INT64 // (N - i): # Overflow

return 0

val *= N - i

val //= i + 1

return val

# Kinda complicated, but works for now

def _register_binom():

# "Fake" UDT so we only compile once for INT64

op = BinaryOp.register_new("binom", _binom, is_udt=True)

typed_op = op[INT64, INT64]

# Make this look like a normal operator

for dtype in [UINT8, UINT16, UINT32, UINT64, INT8, INT16, INT32, INT64]:

op.types[dtype] = INT64

op._typed_ops[dtype] = typed_op

if dtype != INT64:

op.coercions[dtype] = typed_op

# And make it not look like it operates on UDTs

typed_op._type2 = None

op._is_udt = False

op._udt_types = None

op._udt_ops = None

return op

def _first(x, y):

return x # pragma: no cover (numba)

def _second(x, y):

return y # pragma: no cover (numba)

def _pair(x, y):

return 1 # pragma: no cover (numba)

def _udt_dtype(op, dtype, dtype2):

"""Custom dtype handler that compiles an operator for UDTs on first use."""

if dtype._is_udt or dtype2._is_udt:

return op._compile_udt(dtype, dtype2)

def _pair_dtype(op, dtype, dtype2):

return op[INT64]

if _has_numba:

from .udt_utils import (

_compile_codegen,

_has_jit_set,

_maybe_warn_jit_skipped,

compile_udt_binary_wrapper,

set_jit_c_on_op,

)

class BinaryOp(OpBase):

"""Takes two inputs and returns one output, possibly of a different data type.

Built-in and registered BinaryOps are located in the ``graphblas.binary`` namespace

as well as in the ``graphblas.ops`` combined namespace.

"""

__slots__ = (

"_monoid",

"_commutes_to",

"_semiring_commutes_to",

"orig_func",

"is_positional",

"_is_udt",

"_numba_func",

"_custom_dtype",

)

_module = binary

_modname = "binary"

_typed_class = TypedBuiltinBinaryOp

_parse_config = {

"trim_from_front": 4,

"num_underscores": 1,

"re_exprs": [

re.compile(

"^GrB_(FIRST|SECOND|PLUS|MINUS|TIMES|DIV|MIN|MAX)"

"_(BOOL|INT8|UINT8|INT16|UINT16|INT32|UINT32|INT64|UINT64|FP32|FP64|FC32|FC64)$"

),

re.compile(

"GrB_(BOR|BAND|BXOR|BXNOR)_(INT8|INT16|INT32|INT64|UINT8|UINT16|UINT32|UINT64)$"

),

re.compile(

"^GxB_(POW|RMINUS|RDIV|PAIR|ANY|ISEQ|ISNE|ISGT|ISLT|ISGE|ISLE|LOR|LAND|LXOR)"

"_(BOOL|INT8|UINT8|INT16|UINT16|INT32|UINT32|INT64|UINT64|FP32|FP64|FC32|FC64)$"

),

re.compile("^GxB_(FIRST|SECOND|PLUS|MINUS|TIMES|DIV)_(FC32|FC64)$"),

re.compile("^GxB_(ATAN2|HYPOT|FMOD|REMAINDER|LDEXP|COPYSIGN)_(FP32|FP64)$"),

re.compile(

"GxB_(BGET|BSET|BCLR|BSHIFT|FIRSTI1|FIRSTI|FIRSTJ1|FIRSTJ"

"|SECONDI1|SECONDI|SECONDJ1|SECONDJ)"

"_(INT8|INT16|INT32|INT64|UINT8|UINT16|UINT32|UINT64)$"

),

# These are coerced to 0 or 1, but don't return BOOL

re.compile(

"^GxB_(LOR|LAND|LXOR|LXNOR)_"

"(BOOL|INT8|UINT8|INT16|UINT16|INT32|UINT32|INT64|UINT64|FP32|FP64)$"

),

],

"re_exprs_return_bool": [

re.compile("^GrB_(LOR|LAND|LXOR|LXNOR)$"),

re.compile(

"^GrB_(EQ|NE|GT|LT|GE|LE)_"

"(BOOL|INT8|UINT8|INT16|UINT16|INT32|UINT32|INT64|UINT64|FP32|FP64)$"

),

re.compile("^GxB_(EQ|NE)_(FC32|FC64)$"),

],

"re_exprs_return_complex": [re.compile("^GxB_(CMPLX)_(FP32|FP64)$")],

}

_commutes = {

# builtins

"cdiv": "rdiv",

"first": "second",

"ge": "le",

"gt": "lt",

"isge": "isle",

"isgt": "islt",

"minus": "rminus",

"pow": "rpow",

# special

"firsti": "secondi",

"firsti1": "secondi1",

"firstj": "secondj",

"firstj1": "secondj1",

# custom

# "absfirst": "abssecond", # handled in graphblas.binary

# "floordiv": "rfloordiv",

"truediv": "rtruediv",

}

_commutes_to_in_semiring = {

"firsti": "secondj",

"firsti1": "secondj1",

"firstj": "secondi",

"firstj1": "secondi1",

}

_commutative = {

# monoids

"any",

"band",

"bor",

"bxnor",

"bxor",

"eq",

"land",

"lor",

"lxnor",

"lxor",

"max",

"min",

"plus",

"times",

# other

"hypot",

"isclose",

"iseq",

"isne",

"ne",

"pair",

}

# Don't commute: atan2, bclr, bget, bset, bshift, cmplx, copysign, fmod, ldexp, remainder

_positional = {

"firsti",

"firsti1",

"firstj",

"firstj1",

"secondi",

"secondi1",

"secondj",

"secondj1",

}

@classmethod

def _build(cls, name, func, *, is_udt=False, anonymous=False):

if not isinstance(func, FunctionType):

raise TypeError(f"UDF argument must be a function, not {type(func)}")

if name is None:

name = getattr(func, "__name__", "<anonymous_binary>")

success = False

# The error model has to be set here, not only on the cfunc wrapper

# below: ``.compile(sig)`` further down builds the specialization the

# wrapper then reuses, and a Dispatcher keeps one compilation per

# signature. Whichever compile happens first fixes the model, so

# setting it only on the wrapper leaves ``x // 0`` raising

# ZeroDivisionError inside a cfunc, where Numba prints the traceback

# and returns, handing GraphBLAS an element it never wrote.

binary_udf = numba.njit(func, error_model="numpy")

new_type_obj = cls(name, func, anonymous=anonymous, is_udt=is_udt, numba_func=binary_udf)

return_types = {}

nt = numba.types

if not is_udt:

for type_ in _sample_values:

sig = (type_.numba_type, type_.numba_type)

try:

binary_udf.compile(sig)

except numba.TypingError:

continue

ret_type = lookup_dtype(binary_udf.overloads[sig].signature.return_type)

if ret_type != type_ and (

("INT" in ret_type.name and "INT" in type_.name)

or ("FP" in ret_type.name and "FP" in type_.name)

or ("FC" in ret_type.name and "FC" in type_.name)

or (type_ == UINT64 and ret_type == FP64 and return_types.get(INT64) == INT64)

):

# Downcast `ret_type` to `type_`.

# This is what users want most of the time, but we can't make a perfect rule.

# There should be a way for users to be explicit.

ret_type = type_

elif type_ == BOOL and ret_type == INT64 and return_types.get(INT8) == INT8:

ret_type = INT8

input_type = _bool_to_int8(type_)

return_type = _bool_to_int8(ret_type)

# Build wrapper because GraphBLAS wants pointers and void return

wrapper_sig = nt.void(

nt.CPointer(return_type.numba_type),

nt.CPointer(input_type.numba_type),

nt.CPointer(input_type.numba_type),

)

if type_ == BOOL:

if ret_type == BOOL:

def binary_wrapper(z, x, y): # pragma: no cover (numba)

z[0] = bool(binary_udf(bool(x[0]), bool(y[0])))

else:

def binary_wrapper(z, x, y): # pragma: no cover (numba)

z[0] = binary_udf(bool(x[0]), bool(y[0]))

elif ret_type == BOOL:

def binary_wrapper(z, x, y): # pragma: no cover (numba)

z[0] = bool(binary_udf(x[0], y[0]))

else:

def binary_wrapper(z, x, y): # pragma: no cover (numba)

z[0] = binary_udf(x[0], y[0])

binary_wrapper = numba.cfunc(wrapper_sig, nopython=True, error_model="numpy")(

binary_wrapper

)

new_binary = ffi_new("GrB_BinaryOp*")

check_status_carg(

lib.GrB_BinaryOp_new(

new_binary,

binary_wrapper.cffi,

ret_type.gb_obj,

type_.gb_obj,

type_.gb_obj,

),

"BinaryOp",

new_binary[0],

)

op = TypedUserBinaryOp(new_type_obj, name, type_, ret_type, new_binary[0])

new_type_obj._add(op)

success = True

return_types[type_] = ret_type

if success or is_udt:

return new_type_obj

raise UdfParseError("Unable to parse function using Numba")

def _compile_udt(self, dtype, dtype2):

if dtype2 is None:

dtype2 = dtype

dtypes = (dtype, dtype2)

if dtypes in self._udt_types:

return self._udt_ops[dtypes]

# Built-in arithmetic ops set ``_udt_types = {}`` to enable UDT

# dispatch, which also routes plain-scalar misses (e.g. ``plus[BOOL]``)

# here. Re-raise with the legacy single-dtype message so callers like

# ``mapnumpy`` see the same error as before UDT auto-lift.

if self.name in _BUILTIN_UDT_BINARY_OPS and not dtype._is_udt and not dtype2._is_udt:

raise KeyError(f"{self.name} does not work with {dtype}")

if self.name in ("eq", "ne") and not self._anonymous and _has_numba:

binary_wrapper, wrapper_sig = _make_udt_comparison(

dtype, dtype2, is_eq=(self.name == "eq")

)

ret_type = BOOL

elif _has_numba and self.name in _BUILTIN_UDT_BINARY_OPS:

# Auto-generate a per-leaf wrapper for built-in arithmetic. Most

# of these names (plus, minus, times, truediv, min, max) have no

# ``_numba_func`` and would fall through to the KeyError below

# without this branch. ``floordiv`` does have a ``_numba_func``,

# which Numba won't compile against records; the per-leaf wrapper

# is the only path that works for it on UDTs.

py_op = _BUILTIN_UDT_BINARY_OPS[self.name]

binary_wrapper, wrapper_sig, ret_type = compile_udt_binary_wrapper(

self.name, py_op, dtype, dtype2

)

elif self._numba_func is None:

raise KeyError(f"{self.name} does not work with {dtypes} types")

else:

numba_func = self._numba_func

sig = (dtype.numba_type, dtype2.numba_type)

_compile_udf_for_udt(

numba_func, sig, op_kind="binary", op_name=self.name, dtypes=(dtype, dtype2)

)

numba_ret_type = numba_func.overloads[sig].signature.return_type

ret_type = _resolve_udt_return_type(numba_ret_type, dtype, dtype2)

binary_wrapper, wrapper_sig = _get_udt_wrapper(

numba_func, ret_type, dtype, dtype2, numba_ret_type=numba_ret_type

)

op = _finalize_udt_op(

self, dtype, dtype2, ret_type, binary_wrapper, wrapper_sig, TypedUserBinaryOp

)

# Set JIT C definition for auto-generated built-in UDT ops. Only

# same-type pairs apply: mixed UDT+scalar ops don't have a

# meaningful single-type C definition. ``_has_jit_set`` gates access

# to ``lib.GrB_BinaryOp_set_String``, which is absent on SS < 9.

is_jittable = self.name in _BUILTIN_UDT_BINARY_OPS or self.name in ("eq", "ne")

if _has_numba and _has_jit_set and dtype == dtype2 and is_jittable:

if self.name in _BUILTIN_UDT_BINARY_OPS:

op._jit_c_info = set_jit_c_on_op(

op.gb_obj,

self.name,

_BUILTIN_UDT_BINARY_OPS[self.name],

dtype,

lib.GrB_BinaryOp_set_String,

arity=2,

)

else:

# eq/ne return BOOL, not the UDT; the kernel is a chain of

# per-leaf scalar comparisons.

from .udt_utils import set_jit_c_comparison_on_op

op._jit_c_info = set_jit_c_comparison_on_op(

op.gb_obj,

self.name,

dtype,

lib.GrB_BinaryOp_set_String,

is_eq=(self.name == "eq"),

)

_maybe_warn_jit_skipped(op._jit_c_info, self.name, dtype.name)

return op

@classmethod

def register_anonymous(cls, func, name=None, *, parameterized=False, is_udt=False):

"""Register a BinaryOp without registering it in the ``graphblas.binary`` namespace.

Because it is not registered in the namespace, the name is optional.

Parameters

----------

func : FunctionType

The function to compile. For all current backends, this must be able

to be compiled with ``numba.njit``.

``func`` takes two input parameters of any dtype and returns any dtype.

name : str, optional

The name of the operator. This *does not* show up as ``gb.binary.{name}``.

parameterized : bool, default False

When True, create a parameterized user-defined operator, which means

additional parameters can be "baked into" the operator when used.

For example, ``gb.binary.isclose`` is a parameterized function that

optionally accepts ``rel_tol`` and ``abs_tol`` parameters, and it

can be used as: ``A.ewise_mult(B, gb.binary.isclose(rel_tol=1e-5))``.

When creating a parameterized user-defined operator, the ``func``

parameter must be a callable that *returns* a function that will

then get compiled.

is_udt : bool, default False

Whether the operator is intended to operate on user-defined types.

If True, then the function will not be automatically compiled for

builtin types, and it will be compiled "just in time" when used.

Setting ``is_udt=True`` is also helpful when the left and right

dtypes need to be different.

Returns

-------

BinaryOp or ParameterizedBinaryOp

"""

cls._check_supports_udf("register_anonymous")

if parameterized:

return ParameterizedBinaryOp(name, func, anonymous=True, is_udt=is_udt)

return cls._build(name, func, anonymous=True, is_udt=is_udt)

@classmethod

def register_new(cls, name, func, *, parameterized=False, is_udt=False, lazy=False):

"""Register a new BinaryOp and save it to ``graphblas.binary`` namespace.

Parameters

----------

name : str

The name of the operator. This will show up as ``gb.binary.{name}``.

The name may contain periods, ".", which will result in nested objects

such as ``gb.binary.x.y.z`` for name ``"x.y.z"``.

func : FunctionType

The function to compile. For all current backends, this must be able

to be compiled with ``numba.njit``.

``func`` takes two input parameters of any dtype and returns any dtype.

parameterized : bool, default False

When True, create a parameterized user-defined operator, which means

additional parameters can be "baked into" the operator when used.

For example, ``gb.binary.isclose`` is a parameterized function that

optionally accepts ``rel_tol`` and ``abs_tol`` parameters, and it

can be used as: ``A.ewise_mult(B, gb.binary.isclose(rel_tol=1e-5))``.

When creating a parameterized user-defined operator, the ``func``

parameter must be a callable that *returns* a function that will

then get compiled. See the ``user_isclose`` example below.

is_udt : bool, default False

Whether the operator is intended to operate on user-defined types.

If True, then the function will not be automatically compiled for

builtin types, and it will be compiled "just in time" when used.

Setting ``is_udt=True`` is also helpful when the left and right

dtypes need to be different.

lazy : bool, default False

If False (the default), then the function will be automatically

compiled for builtin types (unless ``is_udt`` is True).

Compiling functions can be slow, however, so you may want to

delay compilation and only compile when the operator is used,

which is done by setting ``lazy=True``.

Examples

--------

>>> def max_zero(x, y):

r = 0

if x > r:

r = x

if y > r:

r = y

return r

>>> gb.core.operator.BinaryOp.register_new("max_zero", max_zero)

>>> dir(gb.binary)

[..., 'max_zero', ...]

This is how ``gb.binary.isclose`` is defined:

>>> def user_isclose(rel_tol=1e-7, abs_tol=0.0):

>>> def inner(x, y):

>>> return x == y or abs(x - y) <= max(rel_tol * max(abs(x), abs(y)), abs_tol)

>>> return inner

>>> gb.binary.register_new("user_isclose", user_isclose, parameterized=True)

"""

cls._check_supports_udf("register_new")

module, funcname = cls._remove_nesting(name)

if lazy:

module._delayed[funcname] = (

cls.register_new,

{

"name": name,

"func": func,

"parameterized": parameterized,

"is_udt": is_udt,

},

)

elif parameterized:

binary_op = ParameterizedBinaryOp(name, func, is_udt=is_udt)

setattr(module, funcname, binary_op)

else:

binary_op = cls._build(name, func, is_udt=is_udt)

setattr(module, funcname, binary_op)

# Also save it to `graphblas.op` if not yet defined

opmodule, funcname = cls._remove_nesting(name, module=op, modname="op", strict=False)

if not _hasop(opmodule, funcname):

if lazy:

opmodule._delayed[funcname] = module

else:

setattr(opmodule, funcname, binary_op)

if not cls._initialized:

_STANDARD_OPERATOR_NAMES.add(f"{cls._modname}.{name}")

if not lazy:

return binary_op

@classmethod

def _initialize(cls):

if cls._initialized: # pragma: no cover (safety)

return

super()._initialize()

# Rename div to cdiv

cdiv = binary.cdiv = op.cdiv = BinaryOp("cdiv")

for dtype, ret_type in binary.div.types.items():

orig_op = binary.div[dtype]

cur_op = TypedBuiltinBinaryOp(

cdiv, "cdiv", dtype, ret_type, orig_op.gb_obj, orig_op.gb_name

)

cdiv._add(cur_op)

del binary.div

del op.div

# Add truediv which always points to floating point cdiv

# We are effectively hacking cdiv to always return floating point values

# If the inputs are FP32, we use DIV_FP32; use DIV_FP64 for all other input dtypes

truediv = binary.truediv = op.truediv = BinaryOp("truediv")

rtruediv = binary.rtruediv = op.rtruediv = BinaryOp("rtruediv")

for new_op, builtin_op in [(truediv, binary.cdiv), (rtruediv, binary.rdiv)]:

for dtype in builtin_op.types:

if dtype.name in {"FP32", "FC32", "FC64"}:

orig_dtype = dtype

else:

orig_dtype = FP64

orig_op = builtin_op[orig_dtype]

cur_op = TypedBuiltinBinaryOp(

new_op,

new_op.name,

dtype,

builtin_op.types[orig_dtype],

orig_op.gb_obj,

orig_op.gb_name,

)

new_op._add(cur_op)

if _supports_udfs:

# Add floordiv

# cdiv truncates towards 0, while floordiv truncates towards -inf

BinaryOp.register_new("floordiv", _floordiv, lazy=True) # cast to integer

BinaryOp.register_new("rfloordiv", _rfloordiv, lazy=True) # cast to integer

# For aggregators

BinaryOp.register_new("absfirst", _absfirst, lazy=True)

BinaryOp.register_new("abssecond", _abssecond, lazy=True)

BinaryOp.register_new("rpow", _rpow, lazy=True)

# For algorithms

binary._delayed["binom"] = (_register_binom, {}) # Lazy with custom creation

op._delayed["binom"] = binary

BinaryOp.register_new("isclose", _isclose, parameterized=True)

# Update type information with sane coercion

position_dtypes = [

BOOL,

FP32,

FP64,

INT8,

INT16,

UINT8,

UINT16,

UINT32,

UINT64,

]

if _supports_complex:

position_dtypes.extend([FC32, FC64])

name_types = [

# fmt: off

(

("atan2", "copysign", "fmod", "hypot", "ldexp", "remainder"),

((BOOL, INT8, INT16, UINT8, UINT16), FP32),

((INT32, INT64, UINT32, UINT64), FP64),

),

(

(

"firsti", "firsti1", "firstj", "firstj1", "secondi", "secondi1",

"secondj", "secondj1"),

(

position_dtypes,

INT64,

),

),

(

["lxnor"],

(

(

FP32, FP64, INT8, INT16, INT32, INT64,

UINT8, UINT16, UINT32, UINT64,

),

BOOL,

),

),

# fmt: on

]

if _supports_complex:

name_types.append(

(

["cmplx"],

((BOOL, INT8, INT16, UINT8, UINT16), FP32),

((INT32, INT64, UINT32, UINT64), FP64),

)

)

for names, *types in name_types:

for name in names:

if name in _SS_OPERATORS:

cur_op = binary._deprecated[name]

else:

cur_op = getattr(binary, name)

for input_types, target_type in types:

typed_op = cur_op._typed_ops[target_type]

output_type = cur_op.types[target_type]

for dtype in input_types:

if dtype not in cur_op.types: # pragma: no branch (safety)

cur_op.types[dtype] = output_type

cur_op._typed_ops[dtype] = typed_op

cur_op.coercions[dtype] = target_type

# Not valid input dtypes

del binary.ldexp[FP32]

del binary.ldexp[FP64]

# Fill in commutes info

for left_name, right_name in cls._commutes.items():

if left_name in _SS_OPERATORS:

left = binary._deprecated[left_name]

else:

left = getattr(binary, left_name)

if backend == "suitesparse" and right_name in _SS_OPERATORS:

left._commutes_to = f"ss.{right_name}"

else:

left._commutes_to = right_name

if right_name not in binary._delayed:

if right_name in _SS_OPERATORS:

right = binary._deprecated[right_name]

elif _supports_udfs:

right = getattr(binary, right_name)

else:

right = getattr(binary, right_name, None)

if right is None:

continue

if backend == "suitesparse" and left_name in _SS_OPERATORS:

right._commutes_to = f"ss.{left_name}"

else:

right._commutes_to = left_name

for name in cls._commutative:

if _supports_udfs:

cur_op = getattr(binary, name)

else:

cur_op = getattr(binary, name, None)

if cur_op is None:

continue

cur_op._commutes_to = name

for left_name, right_name in cls._commutes_to_in_semiring.items():

if left_name in _SS_OPERATORS:

left = binary._deprecated[left_name]

else: # pragma: no cover (safety)

left = getattr(binary, left_name)

if right_name in _SS_OPERATORS:

right = binary._deprecated[right_name]

else: # pragma: no cover (safety)

right = getattr(binary, right_name)

left._semiring_commutes_to = right

right._semiring_commutes_to = left

# ``any`` uses ``_second`` so ``monoid.any[udt].reduce(v)`` folds as

# ``acc = v_i`` and returns an actual element. With ``_first`` the

# accumulator would never advance, so the reduce would always return

# the identity.

for binop, func in [

(binary.first, _first),

(binary.second, _second),

(binary.pair, _pair),

(binary.any, _second),

]:

binop.orig_func = func

binop._numba_func = numba.njit(func, error_model="numpy") if _has_numba else None

binop._udt_types = {}

binop._udt_ops = {}

binary.any._numba_func = binary.second._numba_func

binary.first._custom_dtype = _udt_dtype

binary.second._custom_dtype = _udt_dtype

binary.pair._custom_dtype = _pair_dtype

# eq and ne walk the UDT leaf-by-leaf (special-cased in _compile_udt