k

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

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: