Lowering a Reshape to IndexLambda should produce simple expressions if possible

[nkoskelo]

When we lower to an IndexLambda a Reshape expression we output a scalar expression that flattens out the indexing. This is not always necessary. Here is a small example.

import pytato as pt

x = pt.make_placeholder(name="x", shape=(10,4), dtype=float)
expr = pt.expand_dims(x, [2, 3])
idx_lambda = pt.transform.lower_to_index_lambda.to_index_lambda(expr)

first_index_expr = idx_lambda.expr.index_tuple[0]
from pymbolic.primitives import Variable
assert idx_lambda.expr.index_tuple[0] == Variable("_0")

assert idx_lambda.expr.index_tuple[1] == Variable("_1")

Currently, the variable first_index_expr will have structure the ((0 + _0*1) % 10) // 1.

[inducer]

cc @a-alveyblanc

[kaushikcfd]

One option is to use loopy.simplify_using_aff. Although a bit expensive (and potentially not simplifying in the case of SizeParams), leaves us with a straightforward implementation.

[inducer]

In #455, @a-alveyblanc rewired the reshape logic to establish ranges of axes that are "passed through" unchanged. This should require (hopefully) only a simple fix on top of that. I recall there being a discussion at the time about handling of trailing one-long axes, perhaps this just needs to be added.