diff --git a/gem/gem.py b/gem/gem.py index 869daeea..2cf957d0 100644 --- a/gem/gem.py +++ b/gem/gem.py @@ -298,7 +298,8 @@ def __new__(cls, a, b): return a if isinstance(a, Constant) and isinstance(b, Constant): - return Literal(a.value + b.value) + dtype = numpy.result_type(a.dtype, b.dtype) + return Literal(a.value + b.value, dtype=dtype) self = super(Sum, cls).__new__(cls) self.children = a, b @@ -322,7 +323,8 @@ def __new__(cls, a, b): return a if isinstance(a, Constant) and isinstance(b, Constant): - return Literal(a.value * b.value) + dtype = numpy.result_type(a.dtype, b.dtype) + return Literal(a.value * b.value, dtype=dtype) self = super(Product, cls).__new__(cls) self.children = a, b @@ -346,7 +348,8 @@ def __new__(cls, a, b): return a if isinstance(a, Constant) and isinstance(b, Constant): - return Literal(a.value / b.value) + dtype = numpy.result_type(a.dtype, b.dtype) + return Literal(a.value / b.value, dtype=dtype) self = super(Division, cls).__new__(cls) self.children = a, b @@ -410,14 +413,16 @@ def __new__(cls, base, exponent): # Constant folding if isinstance(base, Zero): + dtype = numpy.result_type(base.dtype, exponent.dtype) if isinstance(exponent, Zero): raise ValueError("cannot solve 0^0") - return Zero() + return Zero(dtype=dtype) elif isinstance(exponent, Zero): - return one - - if isinstance(base, Constant) and isinstance(exponent, Constant): - return Literal(base.value ** exponent.value) + dtype = numpy.result_type(base.dtype, exponent.dtype) + return Literal(1, dtype=dtype) + elif isinstance(base, Constant) and isinstance(exponent, Constant): + dtype = numpy.result_type(base.dtype, exponent.dtype) + return Literal(base.value ** exponent.value, dtype=dtype) self = super(Power, cls).__new__(cls) self.children = base, exponent @@ -633,12 +638,12 @@ def __new__(cls, aggregate, multiindex): # Zero folding if isinstance(aggregate, Zero): - return Zero() + return Zero(dtype=aggregate.dtype) # All indices fixed if all(isinstance(i, int) for i in multiindex): if isinstance(aggregate, Constant): - return Literal(aggregate.array[multiindex]) + return Literal(aggregate.array[multiindex], dtype=aggregate.dtype) elif isinstance(aggregate, ListTensor): return aggregate.array[multiindex] @@ -719,7 +724,17 @@ def __init__(self, variable, dim2idxs): idxs_.append((index, stride)) # last += (unknown_extent - 1) * stride elif isinstance(index, int): - offset_ += index * stride + # TODO: Attach dtype to each Node. + # Here, we should simply be able to do: + # >>> offset_ += index * stride + # but "+" and "*" are not currently correctly overloaded + # for indices (integers); they assume floats. + if not isinstance(offset, Integral): + raise NotImplementedError(f"Found non-Integral offset : {offset}") + if isinstance(stride, Constant): + offset_ += index * stride.value + else: + offset_ += index * stride else: raise ValueError("Unexpected index type for flexible indexing") if isinstance(stride, Node): @@ -774,7 +789,7 @@ def __new__(cls, expression, multiindex): # Zero folding if isinstance(expression, Zero): - return Zero(shape) + return Zero(shape, dtype=expression.dtype) self = super(ComponentTensor, cls).__new__(cls) self.children = (expression,) @@ -892,7 +907,8 @@ class Concatenate(Node): def __new__(cls, *children): if all(isinstance(child, Zero) for child in children): size = int(sum(numpy.prod(child.shape, dtype=int) for child in children)) - return Zero((size,)) + dtype = numpy.result_type(*(child.dtype for child in children)) + return Zero((size,), dtype=dtype) self = super(Concatenate, cls).__new__(cls) self.children = children diff --git a/gem/optimise.py b/gem/optimise.py index 7d6c8ecd..7c77cb00 100644 --- a/gem/optimise.py +++ b/gem/optimise.py @@ -15,7 +15,7 @@ Product, Sum, Comparison, Conditional, Division, Index, VariableIndex, Indexed, FlexiblyIndexed, IndexSum, ComponentTensor, ListTensor, Delta, - partial_indexed, one) + partial_indexed, one, uint_type) @singledispatch @@ -182,7 +182,7 @@ def _constant_fold_zero(node, self): def _constant_fold_zero_literal(node, self): if (node.array == 0).all(): # All zeros, make symbolic zero - return Zero(node.shape) + return Zero(node.shape, dtype=node.dtype) else: return node @@ -191,7 +191,8 @@ def _constant_fold_zero_literal(node, self): def _constant_fold_zero_listtensor(node, self): new_children = list(map(self, node.children)) if all(isinstance(nc, Zero) for nc in new_children): - return Zero(node.shape) + dtype = numpy.result_type(*(nc.dtype for nc in new_children)) + return Zero(node.shape, dtype=dtype) elif all(nc == c for nc, c in zip(new_children, node.children)): return node else: @@ -237,7 +238,7 @@ def child(expression): if isinstance(expression, Indexed): return expression.children[0] elif isinstance(expression, Zero): - return Zero(shape) + return Zero(shape, dtype=expression.dtype) return Indexed(_select_expression(list(map(child, expressions)), index), multiindex) if types <= {Literal, Zero, Failure}: @@ -624,7 +625,7 @@ def _replace_delta_delta(node, self): else: def expression(index): if isinstance(index, int): - return Literal(index) + return Literal(index, dtype=uint_type) elif isinstance(index, VariableIndex): return index.expression else: