From 748251fc5f4d6ee3f42294265d896afde68115ea Mon Sep 17 00:00:00 2001 From: Erik Welch Date: Wed, 20 Oct 2021 15:49:00 -0500 Subject: [PATCH] Use numpy utilities in `dtypes.unify`. Also, use numpy behavior when scalars are involved. --- grblas/dtypes.py | 57 ++++++++++++++-------------------------------- grblas/matrix.py | 4 ++-- grblas/operator.py | 24 +++++++++++++++---- grblas/scalar.py | 9 +++++++- grblas/vector.py | 4 ++-- 5 files changed, 49 insertions(+), 49 deletions(-) diff --git a/grblas/dtypes.py b/grblas/dtypes.py index a52689338..845a5a483 100644 --- a/grblas/dtypes.py +++ b/grblas/dtypes.py @@ -1,5 +1,6 @@ import numba import numpy as np +from numpy import find_common_type, promote_types from . import lib @@ -162,7 +163,7 @@ def lookup_dtype(key): raise ValueError(f"Unknown dtype: {key}") -def unify(type1, type2): +def unify(type1, type2, *, is_left_scalar=False, is_right_scalar=False): """ Returns a type that can hold both type1 and type2 @@ -172,43 +173,19 @@ def unify(type1, type2): unify(BOOL, UINT16) -> UINT16 unify(FP32, INT32) -> FP64 """ - if type1 == type2: + if type1 is type2: return type1 - # Bools fit anywhere - if type1 == BOOL: - return type2 - if type2 == BOOL: - return type1 - # Compute bit numbers for comparison - num1 = 8 if type1.name[-1] == "8" else int(type1.name[-2:]) - num2 = 8 if type2.name[-1] == "8" else int(type2.name[-2:]) - - # Same type with different numbers is easy: choose the larger one - if type1.name[1] == type2.name[1]: - return type2 if num2 > num1 else type1 - # One FC, one FP - if type1.name[0] == "F" and type2.name[0] == "F": - return lookup_dtype(f"FC{max(num1, num2)}") - # Float or complex with int - if type1.name[0] == "F": - maxnum = min(64, max(num1, 2 * num2)) - if type1.name[1] == "C": - return lookup_dtype(f"FC{maxnum}") - return lookup_dtype(f"FP{maxnum}") - # Int with float or complex - if type2.name[0] == "F": - maxnum = min(64, max(2 * num1, num2)) - if type2.name[1] == "C": - return lookup_dtype(f"FC{maxnum}") - return lookup_dtype(f"FP{maxnum}") - # Mixed int/uint is a little harder - # Need to double the uint bit number to safely store as int - # Beyond 64, we have to move to float - if type1.name[0] == "U": - num1 *= 2 - else: # type2 is unsigned - num2 *= 2 - maxnum = max(num1, num2) - if maxnum > 64: - return FP64 - return lookup_dtype(f"INT{maxnum}") + if is_left_scalar: + scalar_types = [type1.np_type] + array_types = [] + elif not is_right_scalar: + # Using `promote_types` is faster than `find_common_type` + return lookup_dtype(promote_types(type1.np_type, type2.np_type)) + else: + scalar_types = [] + array_types = [type1.np_type] + if is_right_scalar: + scalar_types.append(type2.np_type) + else: + array_types.append(type2.np_type) + return lookup_dtype(find_common_type(array_types, scalar_types)) diff --git a/grblas/matrix.py b/grblas/matrix.py index e166861e1..3d133b6c6 100644 --- a/grblas/matrix.py +++ b/grblas/matrix.py @@ -566,7 +566,7 @@ def apply(self, op, right=None, *, left=None): keyword_name="left", extra_message="Literal scalars also accepted.", ) - op = get_typed_op(op, self.dtype, left.dtype, kind="binary") + op = get_typed_op(op, left.dtype, self.dtype, is_left_scalar=True, kind="binary") if op.opclass == "Monoid": op = op.binaryop else: @@ -592,7 +592,7 @@ def apply(self, op, right=None, *, left=None): keyword_name="right", extra_message="Literal scalars also accepted.", ) - op = get_typed_op(op, self.dtype, right.dtype, kind="binary") + op = get_typed_op(op, self.dtype, right.dtype, is_right_scalar=True, kind="binary") if op.opclass == "Monoid": op = op.binaryop else: diff --git a/grblas/operator.py b/grblas/operator.py index 7bf0971e3..ff7483894 100644 --- a/grblas/operator.py +++ b/grblas/operator.py @@ -1486,14 +1486,23 @@ def monoid(self): __call__ = TypedBuiltinSemiring.__call__ -def get_typed_op(op, dtype, dtype2=None, *, kind=None): +def get_typed_op(op, dtype, dtype2=None, *, is_left_scalar=False, is_right_scalar=False, kind=None): if isinstance(op, OpBase): if dtype2 is not None: - dtype = unify(dtype, dtype2) + dtype = unify( + dtype, dtype2, is_left_scalar=is_left_scalar, is_right_scalar=is_right_scalar + ) return op[dtype] elif isinstance(op, ParameterizedUdf): op = op() # Use default parameters of parameterized UDFs - return get_typed_op(op, dtype, dtype2, kind=kind) + return get_typed_op( + op, + dtype, + dtype2, + is_left_scalar=is_left_scalar, + is_right_scalar=is_right_scalar, + kind=kind, + ) elif isinstance(op, TypedOpBase): return op elif isinstance(op, Aggregator): @@ -1514,7 +1523,14 @@ def get_typed_op(op, dtype, dtype2=None, *, kind=None): f"Unable to get op from string {op!r}. `kind=` argument must be provided as " '"unary", "binary", "monoid", or "semiring".' ) - return get_typed_op(op, dtype, dtype2, kind=kind) + return get_typed_op( + op, + dtype, + dtype2, + is_left_scalar=is_left_scalar, + is_right_scalar=is_right_scalar, + kind=kind, + ) else: raise TypeError(f"Unable to get typed operator from object with type {type(op)}") diff --git a/grblas/scalar.py b/grblas/scalar.py index 809f1aebe..a7597110c 100644 --- a/grblas/scalar.py +++ b/grblas/scalar.py @@ -144,7 +144,14 @@ def isclose(self, other, *, rel_tol=1e-7, abs_tol=0.0, check_dtype=False): return self.is_empty is other.is_empty # We can't yet call a UDF on a scalar as part of the spec, so let's do it ourselves isclose_func = isclose(rel_tol, abs_tol) - isclose_func = get_typed_op(isclose_func, self.dtype, other.dtype, kind="binary") + isclose_func = get_typed_op( + isclose_func, + self.dtype, + other.dtype, + is_left_scalar=True, + is_right_scalar=True, + kind="binary", + ) return isclose_func.numba_func(self.value, other.value) def clear(self): diff --git a/grblas/vector.py b/grblas/vector.py index 721bba0c4..c5de8e1ad 100644 --- a/grblas/vector.py +++ b/grblas/vector.py @@ -455,7 +455,7 @@ def apply(self, op, right=None, *, left=None): keyword_name="left", extra_message="Literal scalars also accepted.", ) - op = get_typed_op(op, self.dtype, left.dtype, kind="binary") + op = get_typed_op(op, left.dtype, self.dtype, is_left_scalar=True, kind="binary") if op.opclass == "Monoid": op = op.binaryop else: @@ -481,7 +481,7 @@ def apply(self, op, right=None, *, left=None): keyword_name="right", extra_message="Literal scalars also accepted.", ) - op = get_typed_op(op, self.dtype, right.dtype, kind="binary") + op = get_typed_op(op, self.dtype, right.dtype, is_right_scalar=True, kind="binary") if op.opclass == "Monoid": op = op.binaryop else: