Skip to content

Use numpy utilities in dtypes.unify. #123

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
eriknw merged 1 commit into from
Oct 21, 2021
Merged

Use numpy utilities in dtypes.unify. #123

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
57 changes: 17 additions & 40 deletions grblas/dtypes.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
import numba
import numpy as np
from numpy import find_common_type, promote_types

from . import lib

Expand Down Expand Up @@ -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

Expand All @@ -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))
4 changes: 2 additions & 2 deletions grblas/matrix.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down
24 changes: 20 additions & 4 deletions grblas/operator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand All @@ -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)}")

Expand Down
9 changes: 8 additions & 1 deletion grblas/scalar.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
4 changes: 2 additions & 2 deletions grblas/vector.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand All @@ -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:
Expand Down