def ravel(self, shape):
if isinstance(shape, Shape):
nelts = const(1)
for dim in shape.dims:
nelts = self.mul(nelts, dim)
return Shape((nelts, ))
else:
return any_value
def ravel(self, shape):
if isinstance(shape, Shape):
nelts = const(1)
for dim in shape.dims:
nelts = self.mul(nelts, dim)
return Shape((nelts,))
else:
return any_value
def sub(self, x, y):
if is_zero(y):
return x
elif isinstance(x, Const) and isinstance(y, Const):
return const(x.value - y.value)
elif isinstance(x, AnyScalar) or isinstance(y, AnyScalar):
return any_scalar
else:
return Sub(x, y)
def div(self, x, y):
assert not is_zero(y), "Encountered divide by zero during shape inference"
cx = x.__class__
cy = y.__class__
if cx in self._scalar_shape_classes and cy in self._scalar_shape_classes:
if is_one(y):
return x
elif cx is AnyScalar or cy is AnyScalar:
return any_scalar
elif cx is Const and cy is Const:
return const(int(x.value / y.value))
elif x == y:
return const(1)
else:
return Div(x, y)
else:
return any_value
def mul(self, x, y):
if is_zero(x) or is_zero(y):
return const(0)
elif is_one(x):
return y
elif is_one(y):
return x
else:
return Mult(x,y)
def visit_Ravel(self, expr):
shape = self.visit_expr(expr.array)
if isinstance(shape, Shape):
nelts = const(1)
for dim in shape.dims:
nelts = shape_semantics.mul(nelts, dim)
return Shape((nelts,))
else:
return any_value
def div(self, x, y):
assert not is_zero(y)
if is_one(y):
return x
elif isinstance(x, Const) and isinstance(y, Const):
return const(int(x.value / y.value))
elif isinstance(x, AnyScalar) or isinstance(y, AnyScalar):
return any_scalar
else:
return Div(x, y)
def div(self, x, y):
assert not is_zero(
y), "Encountered divide by zero during shape inference"
cx = x.__class__
cy = y.__class__
if cx in self._scalar_shape_classes and cy in self._scalar_shape_classes:
if is_one(y):
return x
elif cx is AnyScalar or cy is AnyScalar:
return any_scalar
elif cx is Const and cy is Const:
return const(int(x.value / y.value))
elif x == y:
return const(1)
else:
return Div(x, y)
else:
return any_value
def add(self, x, y):
if is_zero(x):
return y
elif is_zero(y):
return x
elif isinstance(x, Const) and isinstance(y, Const):
return const(x.value + y.value)
elif isinstance(x, AnyScalar) or isinstance(y, AnyScalar):
return any_scalar
else:
return Add(x,y)
def sub(self, x, y):
cx = x.__class__
cy = y.__class__
if cx in self._scalar_shape_classes and cy in self._scalar_shape_classes:
if is_zero(y):
return x
elif cx is Const and cy is Const:
return const(x.value - y.value)
elif cx is AnyScalar or cy is AnyScalar:
return any_scalar
elif x == y:
return Const(0)
else:
return Sub(x, y)
else:
return any_value
def sub(self, x, y):
cx = x.__class__
cy = y.__class__
if cx in self._scalar_shape_classes and cy in self._scalar_shape_classes:
if is_zero(y):
return x
elif cx is Const and cy is Const:
return const(x.value - y.value)
elif cx is AnyScalar or cy is AnyScalar:
return any_scalar
elif x == y:
return Const(0)
else:
return Sub(x, y)
else:
return any_value
def mul(self, x, y):
cx = x.__class__
cy = y.__class__
if cx in self._scalar_shape_classes and cy in self._scalar_shape_classes:
if is_zero(x) or is_zero(y):
return const(0)
elif is_one(x):
return y
elif is_one(y):
return x
elif cx is AnyScalar or cy is AnyScalar:
return any_scalar
else:
return Mult(x,y)
else:
return any_value
def mul(self, x, y):
cx = x.__class__
cy = y.__class__
if cx in self._scalar_shape_classes and cy in self._scalar_shape_classes:
if is_zero(x) or is_zero(y):
return const(0)
elif is_one(x):
return y
elif is_one(y):
return x
elif cx is AnyScalar or cy is AnyScalar:
return any_scalar
else:
return Mult(x, y)
else:
return any_value
def add(self, x, y):
cx = x.__class__
cy = y.__class__
if cx in self._scalar_shape_classes and cy in self._scalar_shape_classes:
if is_zero(x):
return y
elif is_zero(y):
return x
elif cx is Const and cy is Const:
return const(x.value + y.value)
elif cx is AnyScalar or cy is AnyScalar:
return any_scalar
else:
return Add(x, y)
else:
return any_value
def add(self, x, y):
cx = x.__class__
cy = y.__class__
if cx in self._scalar_shape_classes and cy in self._scalar_shape_classes:
if is_zero(x):
return y
elif is_zero(y):
return x
elif cx is Const and cy is Const:
return const(x.value + y.value)
elif cx is AnyScalar or cy is AnyScalar:
return any_scalar
else:
return Add(x,y)
else:
return any_value
def visit_Array(self, expr): elts = self.visit_expr_list(expr.elts) elt = combine_list(elts) n = len(elts) res = increase_rank(elt, 0, const(n)) return res
def index(self, arr, idx):
if isinstance(arr, Scalar):
return arr
assert arr.__class__ is Shape
if isinstance(idx, (Scalar, Slice)):
indices = [idx]
elif idx.__class__ is Tuple:
indices = idx.elts
else:
assert False, "Unexpected index: %s" % (idx,)
result_dims = []
for (i, curr_idx) in enumerate(indices):
old_dim = arr.dims[i]
if curr_idx is None or \
(isinstance(curr_idx, Const) and curr_idx.value is None):
result_dims.append(old_dim)
elif isinstance(curr_idx, Scalar):
pass
elif curr_idx.__class__ is ConstSlice:
result_dims.append(curr_idx.nelts)
elif curr_idx.__class__ is Shape:
if len(curr_idx.dims) == 0:
# same as unknown scalar
pass
else:
assert len(curr_idx.dims) == 1, "Indexing by a multi-dimensional array not yet supported"
result_dims.append(curr_idx.dims[0])
else:
assert curr_idx.__class__ is Slice, "Unsupported index %s" % curr_idx
if curr_idx.start is None:
lower = const(0)
elif isinstance(curr_idx.start, Const):
if curr_idx.start.value is None:
lower = const(0)
elif curr_idx.start.value < 0:
lower = self.sub(old_dim, curr_idx.start)
else:
lower = curr_idx.start
else:
lower = any_scalar
if curr_idx.stop is None:
upper = old_dim
elif isinstance(curr_idx.stop, Const):
if curr_idx.stop.value is None:
upper = old_dim
elif curr_idx.stop.value < 0:
upper = self.sub(old_dim, curr_idx.stop)
else:
upper = curr_idx.stop
else:
upper = any_scalar
n = self.sub(upper, lower)
step = curr_idx.step
if step and \
isinstance(step, Const) and \
step.value is not None and \
step.value != 1:
n = self.div(n, step)
result_dims.append(n)
n_original = len(arr.dims)
n_idx= len(indices)
if n_original > n_idx:
result_dims.extend(arr.dims[n_idx:])
return make_shape(result_dims)
def visit_Array(self, expr):
elts = self.visit_expr_list(expr.elts)
elt = combine_list(elts)
n = len(elts)
res = increase_rank(elt, 0, const(n))
return res
def int(self, x):
return const(x)
def bool(self, x):
return const(x)
def index(self, arr, idx):
if isinstance(arr, Scalar):
return arr
assert arr.__class__ is Shape
if isinstance(idx, (Scalar, Slice, ConstSlice)):
indices = [idx]
elif idx.__class__ is Tuple:
indices = idx.elts
else:
assert False, "Unexpected index: %s" % (idx, )
result_dims = []
for (i, curr_idx) in enumerate(indices):
old_dim = arr.dims[i]
if curr_idx is None or \
(isinstance(curr_idx, Const) and curr_idx.value is None):
result_dims.append(old_dim)
elif isinstance(curr_idx, Scalar):
pass
elif curr_idx.__class__ is ConstSlice:
result_dims.append(curr_idx.nelts)
elif curr_idx.__class__ is Shape:
if len(curr_idx.dims) == 0:
# same as unknown scalar
pass
else:
assert len(
curr_idx.dims
) == 1, "Indexing by a multi-dimensional array not yet supported"
result_dims.append(curr_idx.dims[0])
else:
assert curr_idx.__class__ is Slice, "Unsupported index %s" % curr_idx
if curr_idx.start is None:
lower = const(0)
elif isinstance(curr_idx.start, Const):
if curr_idx.start.value is None:
lower = const(0)
elif curr_idx.start.value < 0:
lower = self.sub(old_dim, curr_idx.start)
else:
lower = curr_idx.start
else:
lower = any_scalar
if curr_idx.stop is None:
upper = old_dim
elif isinstance(curr_idx.stop, Const):
if curr_idx.stop.value is None:
upper = old_dim
elif curr_idx.stop.value < 0:
upper = self.sub(old_dim, curr_idx.stop)
else:
upper = curr_idx.stop
else:
upper = any_scalar
n = self.sub(upper, lower)
step = curr_idx.step
if step and \
isinstance(step, Const) and \
step.value is not None and \
step.value != 1:
n = self.div(n, step)
result_dims.append(n)
n_original = len(arr.dims)
n_idx = len(indices)
if n_original > n_idx:
result_dims.extend(arr.dims[n_idx:])
return make_shape(result_dims)
def int(self, x): return const(x)
def bool(self, x): return const(x)
