Implements pytato.transform.lower_to_index_lambda.lower_to_index_lambda.

[kaushikcfd]

There's almost no new code here. Moves implementation from PreprocessMapper to fn lower_to_index_lambda(x: Array) -> IndexLambda;.

[inducer]

Do you think we should merge this despite the failing CI?

[kaushikcfd]

I was initially thinking let's wait for #374. But.. it's not downgrading the current situation. So, I would be fine if we merge it.

[inducer]

FWIW, I had neither reviewed nor OK'd this. Sorry if my communication was unclear. I'll take a look now.

[kaushikcfd]

Oh fair. Sorry. Feel free to revert if anything's way out of line here and I'll propose another PR, if it's minor changes I can deal with that right away.

[inducer]

I'm trying to consider the relative merits of your approach here vs. what I proposed.

As a reminder, my proposal was to:

• maybe rename the CodeGenPreprocessor to ToIndexLambdaLowerer or so.
• Implement lower_to_index_lambda by subclassing ToIndexLambdaLowerer and overriding rec to return immediately, effectively making it a single-level, non-recursive affair.

The way I see it, the relative merits are as follows:

This PR:

• lower_to_index_lambda has reasonable cost.
• CodeGenPreprocessor has (I suspect) substantially increased in cost:
  • Dispatch cost is now paid twice, once at the mapper level, and another time at the level of lower_to_index_lambda
  • For each node, first an IndexLambda is created, and that then needs to be re-traversed.
• lower_to_index_lambda is now non-extensible (i.e. customized versions of it cannot be made) because it uses singledispatch.
• lower_to_index_lambda is now inconsistent with the dispatch pattern that's used basically everywhere else in pytato.
• Is a bunch of code churn which was hard to review.

By comparison, the approach I proposed imposes an additional cost of a call of a no-op self.rec for each predecessor on lower_to_index_lambda.

Personally, I can stomach the extra cost on lower_to_index_lambda more easily, but I would value hearing your perspective.

[kaushikcfd]

maybe rename the CodeGenPreprocessor to ToIndexLambdaLowerer or so.

CodeGenPreprocessor does more that lowering to index lambda, like making placeholders from datawrappers, collecting the bound arguments, rename callee kernels). ToIndexLambdaLowerer would be a misnomer.

overriding rec to return immediately, effectively making it a single-level, non-recursive affair.

I did think about that, but then we typically use a Mapper to do a recursive visit over an expression (rec being short for recurse) and rec would need to know whether it was the first call to rec or subsequent ones, otherwise __call__ would just return the expression being traversed. So, I thought a singledispatch was a better tool for the use-case.

lower_to_index_lambda is now non-extensible (i.e. customized versions of it cannot be made) because it uses singledispatch.

Lowering to an index lambda is a precise operation. Unless I'm not thinking quite ahead, I don't see good use cases for it.

[inducer]

ToIndexLambdaLowerer would be a misnomer.

My thought was that bit could be extracted into a base class.

I did think about that, but then we typically use a Mapper to do a recursive visit over an expression (rec being short for recurse) and rec would need to know whether it was the first call to rec or subsequent ones, otherwise __call__ would just return the expression being traversed.

I fail to see the issue here. The outermost dispatch is __call__, which is the inherited rec. The next level is the overridden rec, whose job it is to decide how to do the recursion, namely: not.

So, I thought a singledispatch was a better tool for the use-case.

That part specifically puzzles me. Why go against all existing precedent?

Lowering to an index lambda is a precise operation. Unless I'm not thinking quite ahead, I don't see good use cases for it.

IDK. We're talking about a variant of a preexisting to-index-lambda-lowerer, and you seem to be saying that you're confident that this is the last variant we'll ever need?

Also, want to address the (multiple) cost aspects?

[kaushikcfd]

I fail to see the issue here. The outermost dispatch is call, which is the inherited rec. The next level is the overridden rec, whose job it is to decide how to do the recursion, namely: not.

Let me give this a shot. I'm not entirely sure if it will work, but I will propose a PR soon and we can iterate over there.

Also, want to address the (multiple) cost aspects?

That I concede. :)