def facet_integral_predicates(self, mesh, integral_type, kinfo): self.bag.needs_cell_facets = True # Number of recerence cell facets if mesh.cell_set._extruded: self.num_facets = mesh._base_mesh.ufl_cell().num_facets() else: self.num_facets = mesh.ufl_cell().num_facets() # Index for loop over cell faces of reference cell fidx = self.bag.index_creator((self.num_facets,)) # Cell is interior or exterior select = 1 if integral_type.startswith("interior_facet") else 0 i = self.bag.index_creator((1,)) predicates = [pym.Comparison(pym.Subscript(pym.Variable(self.cell_facets_arg), (fidx[0], 0)), "==", select)] # TODO subdomain boundary integrals, this does the wrong thing for integrals like f*ds + g*ds(1) # "otherwise" is treated incorrectly as "everywhere" # However, this replicates an existing slate bug. if kinfo.subdomain_id != "otherwise": predicates.append(pym.Comparison(pym.Subscript(pym.Variable(self.cell_facets_arg), (fidx[0], 1)), "==", kinfo.subdomain_id)) # Additional facet array argument to be fed into tsfc loopy kernel subscript = pym.Subscript(pym.Variable(self.local_facet_array_arg), (pym.Sum((i[0], fidx[0])))) facet_arg = SubArrayRef(i, subscript) return predicates, fidx, facet_arg
def expression_indexed(expr, parameters): aggregate, multiindex = (expression(c, parameters) for c in expr.children) return pym.Subscript(aggregate, multiindex) extents = [int(numpy.prod(expr.aggregate.shape[i+1:])) for i in range(len(multiindex))] make_sum = lambda x, y: pym.Sum((x, y)) index = reduce(make_sum, [pym.Product((e, m)) for e, m in zip(extents, multiindex)]) return pym.Subscript(aggregate, (index,))
def map_non_contiguous_advanced_index(self, expr: AdvancedIndexInNoncontiguousAxes ) -> IndexLambda: from pytato.utils import (get_shape_after_broadcasting, get_indexing_expression) i_adv_indices = tuple(i for i, idx_expr in enumerate(expr.indices) if isinstance(idx_expr, (Array, INT_CLASSES))) adv_idx_shape = get_shape_after_broadcasting([expr.indices[i_idx] for i_idx in i_adv_indices]) vng = UniqueNameGenerator() indices = [] in_ary = vng("in") bindings = {in_ary: self.rec(expr.array)} islice_idx = len(adv_idx_shape) for idx, axis_len in zip(expr.indices, expr.array.shape): if isinstance(idx, INT_CLASSES): if isinstance(axis_len, INT_CLASSES): indices.append(idx % axis_len) else: bnd_name = vng("in") bindings[bnd_name] = self.rec(axis_len) indices.append(idx % prim.Variable(bnd_name)) elif isinstance(idx, NormalizedSlice): indices.append(idx.start + idx.step * prim.Variable(f"_{islice_idx}")) islice_idx += 1 elif isinstance(idx, Array): if isinstance(axis_len, INT_CLASSES): bnd_name = vng("in") bindings[bnd_name] = self.rec(idx) indirect_idx_expr = prim.Subscript(prim.Variable(bnd_name), get_indexing_expression( idx.shape, adv_idx_shape)) if not idx.tags_of_type(AssumeNonNegative): indirect_idx_expr = indirect_idx_expr % axis_len indices.append(indirect_idx_expr) else: raise NotImplementedError("Advanced indexing over" " parametric axis lengths.") else: raise NotImplementedError(f"Indices of type {type(idx)}.") return IndexLambda(expr=prim.Subscript(prim.Variable(in_ary), tuple(indices)), bindings=bindings, shape=expr.shape, dtype=expr.dtype, axes=expr.axes, tags=expr.tags, )
def initialise_terminals(self, var2terminal, coefficients): """ Initilisation of the variables in which coefficients and the Tensors coming from TSFC are saved. :arg var2terminal: dictionary that maps Slate Tensors to gem Variables """ tensor2temp = OrderedDict() inits = [] for gem_tensor, slate_tensor in var2terminal.items(): assert slate_tensor.terminal, "Only terminal tensors need to be initialised in Slate kernels." (_, dtype), = assign_dtypes([gem_tensor], self.tsfc_parameters["scalar_type"]) loopy_tensor = loopy.TemporaryVariable( gem_tensor.name, dtype=dtype, shape=gem_tensor.shape, address_space=loopy.AddressSpace.LOCAL) tensor2temp[slate_tensor] = loopy_tensor if not slate_tensor.assembled: indices = self.bag.index_creator(self.shape(slate_tensor)) inames = {var.name for var in indices} var = pym.Subscript(pym.Variable(loopy_tensor.name), indices) inits.append( loopy.Assignment(var, "0.", id="init%d" % len(inits), within_inames=frozenset(inames))) else: f = slate_tensor.form if isinstance( slate_tensor.form, tuple) else (slate_tensor.form, ) coeff = tuple(coefficients[c] for c in f) offset = 0 ismixed = tuple( (type(c.ufl_element()) == MixedElement) for c in f) names = [] for (im, c) in zip(ismixed, coeff): names += [name for (name, ext) in c.values()] if im else [c[0]] # Mixed coefficients come as seperate parameter (one per space) for i, shp in enumerate(*slate_tensor.shapes.values()): indices = self.bag.index_creator((shp, )) inames = {var.name for var in indices} offset_index = (pym.Sum((offset, indices[0])), ) name = names[i] if ismixed else names var = pym.Subscript(pym.Variable(loopy_tensor.name), offset_index) c = pym.Subscript(pym.Variable(name), indices) inits.append( loopy.Assignment(var, c, id="init%d" % len(inits), within_inames=frozenset(inames))) offset += shp return inits, tensor2temp
def map_basic_index(self, expr: BasicIndex) -> IndexLambda: vng = UniqueNameGenerator() indices = [] in_ary = vng("in") bindings = {in_ary: self.rec(expr.array)} islice_idx = 0 for idx, axis_len in zip(expr.indices, expr.array.shape): if isinstance(idx, INT_CLASSES): if isinstance(axis_len, INT_CLASSES): indices.append(idx % axis_len) else: bnd_name = vng("in") bindings[bnd_name] = axis_len indices.append(idx % prim.Variable(bnd_name)) elif isinstance(idx, NormalizedSlice): indices.append(idx.start + idx.step * prim.Variable(f"_{islice_idx}")) islice_idx += 1 else: raise NotImplementedError return IndexLambda(expr=prim.Subscript(prim.Variable(in_ary), tuple(indices)), bindings=bindings, shape=expr.shape, dtype=expr.dtype, axes=expr.axes, tags=expr.tags, )
def build_ass(): # A_T[i,j] = sum(k, A0[i,j,k] * G_T[k]); # Get variable symbols for all required variables i, j, k = inames["i"], inames["j"], inames["k"] A_T, A0, G_T = args["A_T"], args["A0"], args["G_T"] # The target of the assignment target = pb.Subscript(A_T, (i, j)) # The rhs expression: Frobenius inner product <A0[i,j],G_T> reduce_op = lp.library.reduction.SumReductionOperation() reduce_expr = pb.Subscript(A0, (i, j, k)) * pb.Subscript(G_T, (k)) expr = lp.Reduction(reduce_op, k, reduce_expr) return lp.Assignment(target, expr)
def pymbolic_variable(self, node): pym = self._gem_to_pym_var(node) if node in self.indices: indices = self.fetch_multiindex(self.indices[node]) if indices: return p.Subscript(pym, indices) return pym
def generate_lhs(self, tensor, temp): """ Generation of an lhs for the loopy kernel, which contains the TSFC assembly of the tensor. """ idx = self.bag.index_creator(self.shape(tensor)) lhs = pym.Subscript(temp, idx) return SubArrayRef(idx, lhs)
def _expression_indexed(expr, ctx): rank = ctx.pym_multiindex(expr.multiindex) var = expression(expr.children[0], ctx) if isinstance(var, p.Subscript): rank = var.index + rank var = var.aggregate return p.Subscript(var, rank)
def _apply_elem_wise_func(x: Array, func_name: str) -> IndexLambda: if x.dtype.kind != "f": raise ValueError(f"'{func_name}' does not support '{x.dtype}' arrays.") expr = prim.Call(var(f"pytato.c99.{func_name}"), (prim.Subscript( var("in"), tuple(var(f"_{i}") for i in range(len(x.shape)))), )) return IndexLambda(x.namespace, expr, x.shape, x.dtype, {"in": x})
def map_SympyField(self, expr): f = expr.field if expr.subscript is not None: subscript = tuple(self.rec(i) for i in expr.subscript) return pp.Subscript(f, subscript) else: return f
def build_ass(): """ A[i,j] = c*sum(k, B[k,i]*B[k,j]) """ # The target of the assignment target = pb.Subscript(args["A"], (inames["i"], inames["j"])) # The rhs expression: A reduce operation of the matrix columns # Maybe replace with manual increment? reduce_op = lp.library.reduction.SumReductionOperation() reduce_expr = pb.Subscript(args["B"], (inames["k"], inames["i"])) * pb.Subscript( args["B"], (inames["k"], inames["j"])) expr = args["c"] * lp.Reduction(reduce_op, inames["k"], reduce_expr) return lp.Assignment(target, expr)
def test_math_function(target, tp): # Test correct maths functions are generated for C and OpenCL # backend instead for different data type data_type = {"f32": np.float32, "f64": np.float64}[tp] import pymbolic.primitives as p i = p.Variable("i") xi = p.Subscript(p.Variable("x"), i) yi = p.Subscript(p.Variable("y"), i) zi = p.Subscript(p.Variable("z"), i) n = 100 domain = "{[i]: 0<=i<%d}" % n data = [ lp.GlobalArg("x", data_type, shape=(n, )), lp.GlobalArg("y", data_type, shape=(n, )), lp.GlobalArg("z", data_type, shape=(n, )) ] inst = [lp.Assignment(xi, p.Variable("min")(yi, zi))] knl = lp.make_kernel(domain, inst, data, target=target()) code = lp.generate_code_v2(knl).device_code() assert "fmin" in code if tp == "f32" and target == CTarget: assert "fminf" in code else: assert "fminf" not in code inst = [lp.Assignment(xi, p.Variable("max")(yi, zi))] knl = lp.make_kernel(domain, inst, data, target=target()) code = lp.generate_code_v2(knl).device_code() assert "fmax" in code if tp == "f32" and target == CTarget: assert "fmaxf" in code else: assert "fmaxf" not in code
def initialise_terminals(self, var2terminal, coefficients): """ Initilisation of the variables in which coefficients and the Tensors coming from TSFC are saved. :arg var2terminal: dictionary that maps Slate Tensors to gem Variables """ tensor2temp = OrderedDict() inits = [] for gem_tensor, slate_tensor in var2terminal.items(): loopy_tensor = loopy.TemporaryVariable(gem_tensor.name, shape=gem_tensor.shape, address_space=loopy.AddressSpace.LOCAL) tensor2temp[slate_tensor] = loopy_tensor if isinstance(slate_tensor, slate.Tensor): indices = self.bag.index_creator(self.shape(slate_tensor)) inames = {var.name for var in indices} var = pym.Subscript(pym.Variable(loopy_tensor.name), indices) inits.append(loopy.Assignment(var, "0.", id="init%d" % len(inits), within_inames=frozenset(inames))) elif isinstance(slate_tensor, slate.AssembledVector): f = slate_tensor._function coeff = coefficients[f] offset = 0 ismixed = (type(f.ufl_element()) == MixedElement) names = [name for (name, ext) in coeff.values()] if ismixed else coeff[0] # Mixed coefficients come as seperate parameter (one per space) for i, shp in enumerate(*slate_tensor.shapes.values()): indices = self.bag.index_creator((shp,)) inames = {var.name for var in indices} offset_index = (pym.Sum((offset, indices[0])),) name = names[i] if ismixed else names var = pym.Subscript(pym.Variable(loopy_tensor.name), offset_index) c = pym.Subscript(pym.Variable(name), indices) inits.append(loopy.Assignment(var, c, id="init%d" % len(inits), within_inames=frozenset(inames))) offset += shp return inits, tensor2temp
def map_field(self, expr, *args, **kwargs): if expr.ignore_prepends: pre_index = () else: prepend = kwargs.get("prepend_with") or () pre_index = tuple(parse_if_str(x) for x in prepend) pre_index = pre_index + kwargs.pop("outer_subscript", ()) full_index = pre_index + expr.index_tuple if full_index == tuple(): x = expr.child else: if isinstance(expr.child, pp.Subscript): full_index = pre_index + expr.child.index_tuple + expr.index_tuple x = pp.Subscript(expr.child.aggregate, self.rec(full_index)) else: x = pp.Subscript(expr.child, self.rec(full_index)) return self.rec(x, *args, **kwargs)
def loopify_tsfc_kernel_data(self, kernel_data): """ This method generates loopy arguments from the kernel data, which are then fed to the TSFC loopy kernel. The arguments are arrays and have to be fed element by element to loopy aka they have to be subarrayrefed. """ arguments = [] for c, name in kernel_data: extent = self.extent(c) idx = self.bag.index_creator(extent) arguments.append(SubArrayRef(idx, pym.Subscript(pym.Variable(name), idx))) return arguments
def statement_evaluate(leaf, ctx): expr = leaf.expression if isinstance(expr, gem.ListTensor): ops = [] var, index = ctx.pymbolic_variable_and_destruct(expr) for multiindex, value in numpy.ndenumerate(expr.array): ops.append( lp.Assignment(p.Subscript(var, index + multiindex), expression(value, ctx), within_inames=ctx.active_inames())) return ops elif isinstance(expr, gem.Constant): return [] elif isinstance(expr, gem.ComponentTensor): idx = ctx.gem_to_pym_multiindex(expr.multiindex) var, sub_idx = ctx.pymbolic_variable_and_destruct(expr) lhs = p.Subscript(var, idx + sub_idx) with active_indices(dict(zip(expr.multiindex, idx)), ctx) as ctx_active: return [ lp.Assignment(lhs, expression(expr.children[0], ctx_active), within_inames=ctx_active.active_inames()) ] elif isinstance(expr, gem.Inverse): idx = ctx.pymbolic_multiindex(expr.shape) var = ctx.pymbolic_variable(expr) lhs = (SubArrayRef(idx, p.Subscript(var, idx)), ) idx_reads = ctx.pymbolic_multiindex(expr.children[0].shape) var_reads = ctx.pymbolic_variable(expr.children[0]) reads = (SubArrayRef(idx_reads, p.Subscript(var_reads, idx_reads)), ) rhs = p.Call(p.Variable("inverse"), reads) return [ lp.CallInstruction(lhs, rhs, within_inames=ctx.active_inames()) ] elif isinstance(expr, gem.Solve): idx = ctx.pymbolic_multiindex(expr.shape) var = ctx.pymbolic_variable(expr) lhs = (SubArrayRef(idx, p.Subscript(var, idx)), ) reads = [] for child in expr.children: idx_reads = ctx.pymbolic_multiindex(child.shape) var_reads = ctx.pymbolic_variable(child) reads.append( SubArrayRef(idx_reads, p.Subscript(var_reads, idx_reads))) rhs = p.Call(p.Variable("solve"), tuple(reads)) return [ lp.CallInstruction(lhs, rhs, within_inames=ctx.active_inames()) ] else: return [ lp.Assignment(ctx.pymbolic_variable(expr), expression(expr, ctx, top=True), within_inames=ctx.active_inames()) ]
def slate_call(self, kernel, temporaries): # Slate kernel call reads = [] for t in temporaries: shape = t.shape name = t.name idx = self.bag.index_creator(shape) reads.append(SubArrayRef(idx, pym.Subscript(pym.Variable(name), idx))) call = pym.Call(pym.Variable(kernel.name), tuple(reads)) output_var = pym.Variable(kernel.args[0].name) slate_kernel_call_output = self.generate_lhs(self.expression, output_var) insn = loopy.CallInstruction((slate_kernel_call_output,), call, id="slate_kernel_call") return insn
def _get_sub_array_ref(array: Array, name: str) -> "lp.symbolic.SubArrayRef": inames = tuple( state.var_name_gen(f"_{name}_dim{d}") for d in range(array.ndim)) domains.append( domain_for_shape( inames, shape_to_scalar_expression(array.shape, self, state), {})) inames_as_vars = tuple(var(iname) for iname in inames) return SubArrayRef(inames_as_vars, prim.Subscript(var(name), inames_as_vars))
def _expression_flexiblyindexed(expr, ctx): var = expression(expr.children[0], ctx) rank = [] for off, idxs in expr.dim2idxs: for index, stride in idxs: assert isinstance(index, gem.Index) rank_ = [off] for index, stride in idxs: rank_.append(p.Product((ctx.active_indices[index], stride))) rank.append(p.Sum(tuple(rank_))) return p.Subscript(var, tuple(rank))
def pymbolic_variable(self, node): try: pym = self.gem_to_pymbolic[node] except KeyError: name = self.name_gen(node.name) pym = p.Variable(name) self.gem_to_pymbolic[node] = pym if node in self.indices: indices = self.pym_multiindex(self.indices[node]) if indices: return p.Subscript(pym, indices) else: return pym else: return pym
def statement_evaluate(leaf, ctx): expr = leaf.expression if isinstance(expr, gem.ListTensor): ops = [] var = ctx.pymbolic_variable(expr) index = () if isinstance(var, p.Subscript): var, index = var.aggregate, var.index_tuple for multiindex, value in numpy.ndenumerate(expr.array): ops.append(lp.Assignment(p.Subscript(var, index + multiindex), expression(value, ctx), within_inames=ctx.active_inames())) return ops elif isinstance(expr, gem.Constant): return [] else: return [lp.Assignment(ctx.pymbolic_variable(expr), expression(expr, ctx, top=True), within_inames=ctx.active_inames())]
def _apply_elem_wise_func(inputs: Tuple[ArrayOrScalar, ...], func_name: str, ret_dtype: Optional[_dtype_any] = None, np_func_name: Optional[str] = None ) -> ArrayOrScalar: if all(isinstance(x, SCALAR_CLASSES) for x in inputs): if np_func_name is None: np_func_name = func_name np_func = getattr(np, np_func_name) return np_func(*inputs) # type: ignore if not inputs: raise ValueError("at least one argument must be present") shape = None sym_args = [] bindings = {} for index, inp in enumerate(inputs): if isinstance(inp, Array): if inp.dtype.kind not in ["f", "c"]: raise ValueError("only floating-point or complex " "arguments supported") if shape is None: shape = inp.shape elif inp.shape != shape: # FIXME: merge this logic with arithmetic, so that broadcasting # is implemented properly raise NotImplementedError("broadcasting in function application") if ret_dtype is None: ret_dtype = inp.dtype bindings[f"in_{index}"] = inp sym_args.append( prim.Subscript(var(f"in_{index}"), tuple(var(f"_{i}") for i in range(len(shape))))) else: sym_args.append(inp) assert shape is not None assert ret_dtype is not None return IndexLambda( prim.Call(var(f"pytato.c99.{func_name}"), tuple(sym_args)), shape, ret_dtype, bindings, axes=_get_default_axes(len(shape)))
def map_Subscript(self, expr): # noqa # (expr value, slice slice, expr_context ctx) def none_or_rec(x): if x is None: return x else: return self.rec(x) if isinstance(expr.slice, slice): index = slice(none_or_rec(expr.slice.start), none_or_rec(expr.slice.stop), none_or_rec(expr.slice.step)) else: index = none_or_rec(expr.slice) return p.Subscript(self.rec(expr.value), index)
def test_multi_arg_array_call(ctx_factory): ctx = ctx_factory() queue = cl.CommandQueue(ctx) import pymbolic.primitives as p n = 10 acc_i = p.Variable("acc_i") i = p.Variable("i") index = p.Variable("index") a_i = p.Subscript(p.Variable("a"), p.Variable("i")) argmin_kernel = lp.make_function("{[i]: 0 <= i < n}", [ lp.Assignment(id="init2", assignee=index, expression=0), lp.Assignment(id="init1", assignee=acc_i, expression="214748367"), lp.Assignment(id="insn", assignee=index, expression=p.If(p.Expression.eq(acc_i, a_i), i, index), depends_on="update"), lp.Assignment(id="update", assignee=acc_i, expression=p.Variable("min")(acc_i, a_i), depends_on="init1,init2") ], [ lp.GlobalArg("a"), lp.GlobalArg( "acc_i, index", is_input=False, is_output=True, shape=lp.auto), ... ], name="custom_argmin") argmin_kernel = lp.fix_parameters(argmin_kernel, n=n) knl = lp.make_kernel( "{[i]:0<=i<n}", """ []: min_val[()], []: min_index[()] = custom_argmin([i]:b[i]) """) knl = lp.fix_parameters(knl, n=n) knl = lp.set_options(knl, return_dict=True) knl = lp.merge([knl, argmin_kernel]) b = np.random.randn(n) evt, out_dict = knl(queue, b=b) tol = 1e-15 from numpy.linalg import norm assert (norm(out_dict["min_val"] - np.min(b)) < tol) assert (norm(out_dict["min_index"] - np.argmin(b)) < tol)
def _make_reduction_lambda(op: ReductionOperation, a: Array, axis: Optional[Union[int, Tuple[int]]], initial: Any) -> Array: """ Return a :class:`IndexLambda` that performs reduction over the *axis* axes of *a* with the reduction op *op*. :arg op: The reduction operation to perform. :arg a: The :class:`pytato.Array` on which to perform the reduction. :arg axis: The axes over which the reduction is to be performed. If axis is *None*, perform reduction over all of *a*'s axes. """ new_shape, reduction_axes = _normalize_reduction_axes(a.shape, axis) del axis indices, redn_bounds = _get_reduction_indices_bounds( a.shape, reduction_axes) if initial is _NoValue: for iax in reduction_axes: shape_iax = a.shape[iax] from pytato.utils import are_shape_components_equal if are_shape_components_equal(shape_iax, 0): raise ValueError( "zero-size reduction operation with no supplied " "'initial' value") if isinstance(iax, Array): raise NotImplementedError( "cannot statically determine emptiness of " f"reduction axis {iax} (0-based)") elif initial != op.neutral_element(a.dtype): raise NotImplementedError("reduction with 'initial' not equal to the " "neutral element") return make_index_lambda( Reduce(prim.Subscript(prim.Variable("in"), tuple(indices)), op, redn_bounds), {"in": a}, new_shape, a.dtype)
def _as_array_or_scalar(exprs: Sequence[ScalarExpression], bindings: Mapping[str, Array], out_shape: ShapeType ) -> Tuple[ArrayOrScalar, ...]: """ Helper routine invoked in :func:`index_lambda_to_high_level_op`. For every expression in *exprs* either infers (and returns) it as a scalar or infers (and returns) the array corresponding to one of *bindings* while defining an :class:`~pytato.array.IndexLambda` of *out_shape* shape. """ result: List[ArrayOrScalar] = [] if out_shape != get_shape_after_broadcasting(bindings.values()): raise UnknownIndexLambdaExpr() binding_to_subscript = {bnd_name: p.Subscript( p.Variable(bnd_name), get_indexing_expression(bnd.shape, out_shape)) for bnd_name, bnd in bindings.items()} for expr in exprs: if isinstance(expr, SCALAR_CLASSES): result.append(expr) elif (isinstance(expr, p.Variable) and bindings[expr.name].shape == ()): result.append(bindings[expr.name]) elif (isinstance(expr, p.Subscript) and isinstance(expr.aggregate, p.Variable) and (binding_to_subscript[expr.aggregate.name] == expr)): result.append(bindings[expr.aggregate.name]) else: raise UnknownIndexLambdaExpr() return tuple(result)
def map_einsum(self, expr: Einsum) -> Array: import operator from functools import reduce from pytato.scalar_expr import Reduce from pytato.utils import (dim_to_index_lambda_components, are_shape_components_equal) from pytato.array import ElementwiseAxis, ReductionAxis bindings = {f"in{k}": self.rec(arg) for k, arg in enumerate(expr.args)} redn_bounds: Dict[str, Tuple[ScalarExpression, ScalarExpression]] = {} args_as_pym_expr: List[prim.Subscript] = [] namegen = UniqueNameGenerator(set(bindings)) # {{{ add bindings coming from the shape expressions for access_descr, (iarg, arg) in zip(expr.access_descriptors, enumerate(expr.args)): subscript_indices = [] for iaxis, axis in enumerate(access_descr): if not are_shape_components_equal( arg.shape[iaxis], expr._access_descr_to_axis_len()[axis]): # axis is broadcasted assert are_shape_components_equal(arg.shape[iaxis], 1) subscript_indices.append(0) continue if isinstance(axis, ElementwiseAxis): subscript_indices.append(prim.Variable(f"_{axis.dim}")) else: assert isinstance(axis, ReductionAxis) redn_idx_name = f"_r{axis.dim}" if redn_idx_name not in redn_bounds: # convert the ShapeComponent to a ScalarExpression redn_bound, redn_bound_bindings = ( dim_to_index_lambda_components( arg.shape[iaxis], namegen)) redn_bounds[redn_idx_name] = (0, redn_bound) bindings.update({k: self.rec(v) for k, v in redn_bound_bindings.items()}) subscript_indices.append(prim.Variable(redn_idx_name)) args_as_pym_expr.append(prim.Subscript(prim.Variable(f"in{iarg}"), tuple(subscript_indices))) # }}} inner_expr = reduce(operator.mul, args_as_pym_expr[1:], args_as_pym_expr[0]) if redn_bounds: from pytato.reductions import SumReductionOperation inner_expr = Reduce(inner_expr, SumReductionOperation(), redn_bounds) return IndexLambda(expr=inner_expr, shape=tuple(self.rec(s) if isinstance(s, Array) else s for s in expr.shape), dtype=expr.dtype, bindings=bindings, axes=expr.axes, tags=expr.tags)
def parse_postfix(self, pstate, min_precedence, left_exp): import pymbolic.primitives as primitives import pymbolic.parser as p did_something = False next_tag = pstate.next_tag() if next_tag is p._openpar and p._PREC_CALL > min_precedence: pstate.advance() pstate.expect_not_end() if next_tag is p._closepar: pstate.advance() left_exp = primitives.Call(left_exp, ()) else: args = self.parse_expression(pstate) if not isinstance(args, tuple): args = (args, ) pstate.expect(p._closepar) pstate.advance() if left_exp == primitives.Variable("matrix"): left_exp = np.array(list(list(row) for row in args)) else: left_exp = primitives.Call(left_exp, args) did_something = True elif next_tag is p._openbracket and p._PREC_CALL > min_precedence: pstate.advance() pstate.expect_not_end() left_exp = primitives.Subscript(left_exp, self.parse_expression(pstate)) pstate.expect(p._closebracket) pstate.advance() did_something = True elif next_tag is p._dot and p._PREC_CALL > min_precedence: pstate.advance() pstate.expect(p._identifier) left_exp = primitives.Lookup(left_exp, pstate.next_str()) pstate.advance() did_something = True elif next_tag is p._plus and p._PREC_PLUS > min_precedence: pstate.advance() left_exp += self.parse_expression(pstate, p._PREC_PLUS) did_something = True elif next_tag is p._minus and p._PREC_PLUS > min_precedence: pstate.advance() left_exp -= self.parse_expression(pstate, p._PREC_PLUS) did_something = True elif next_tag is p._times and p._PREC_TIMES > min_precedence: pstate.advance() left_exp *= self.parse_expression(pstate, p._PREC_TIMES) did_something = True elif next_tag is p._over and p._PREC_TIMES > min_precedence: pstate.advance() from pymbolic.primitives import Quotient left_exp = Quotient(left_exp, self.parse_expression(pstate, p._PREC_TIMES)) did_something = True elif next_tag is self.power_sym and p._PREC_POWER > min_precedence: pstate.advance() exponent = self.parse_expression(pstate, p._PREC_POWER) if left_exp == np.e: from pymbolic.primitives import Call, Variable left_exp = Call(Variable("exp"), (exponent, )) else: left_exp **= exponent did_something = True elif next_tag is p._comma and p._PREC_COMMA > min_precedence: # The precedence makes the comma left-associative. pstate.advance() if pstate.is_at_end() or pstate.next_tag() is p._closepar: left_exp = (left_exp, ) else: new_el = self.parse_expression(pstate, p._PREC_COMMA) if isinstance(left_exp, tuple) \ and not isinstance(left_exp, FinalizedTuple): left_exp = left_exp + (new_el, ) else: left_exp = (left_exp, new_el) did_something = True return left_exp, did_something
def map_Indexed(self, expr): # noqa return prim.Subscript( self.rec(expr.args[0].args[0]), tuple(self.rec(i) for i in expr.args[1:]) )