diff --git a/.basedpyright/baseline.json b/.basedpyright/baseline.json index dff7a32a5..c5cbc4be0 100644 --- a/.basedpyright/baseline.json +++ b/.basedpyright/baseline.json @@ -7623,46 +7623,6 @@ "lineCount": 1 } }, - { - "code": "reportReturnType", - "range": { - "startColumn": 11, - "endColumn": 87, - "lineCount": 4 - } - }, - { - "code": "reportUnknownVariableType", - "range": { - "startColumn": 11, - "endColumn": 87, - "lineCount": 4 - } - }, - { - "code": "reportOperatorIssue", - "range": { - "startColumn": 8, - "endColumn": 46, - "lineCount": 1 - } - }, - { - "code": "reportOperatorIssue", - "range": { - "startColumn": 8, - "endColumn": 60, - "lineCount": 1 - } - }, - { - "code": "reportUnknownArgumentType", - "range": { - "startColumn": 8, - "endColumn": 86, - "lineCount": 2 - } - }, { "code": "reportUnknownMemberType", "range": { diff --git a/.gitignore b/.gitignore index 455360283..4b745f419 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ distribute*tar.gz doc/_build .mypy_cache +.venv *.dot *.svg diff --git a/pytato/scalar_expr.py b/pytato/scalar_expr.py index f99c2280c..06219e540 100644 --- a/pytato/scalar_expr.py +++ b/pytato/scalar_expr.py @@ -73,6 +73,7 @@ ) from pymbolic.mapper.distributor import DistributeMapper as DistributeMapperBase from pymbolic.mapper.evaluator import EvaluationMapper as EvaluationMapperBase +from pymbolic.mapper.flattener import FlattenMapper as FlattenMapperBase from pymbolic.mapper.stringifier import StringifyMapper as StringifyMapperBase from pymbolic.mapper.substitutor import SubstitutionMapper as SubstitutionMapperBase from pymbolic.typing import Integer @@ -212,18 +213,20 @@ def map_reduce(self, expr: Reduce) -> Never: class DistributeMapper(DistributeMapperBase): - def map_reduce(self, expr: Reduce) -> None: # TODO: not trivial to distribute symbolic reduction nodes raise NotImplementedError() class TermCollector(TermCollectorBase): - def map_reduce(self, expr: Reduce) -> None: raise NotImplementedError() +class FlattenMapper(FlattenMapperBase, IdentityMapper[[]]): + pass + + class StringifyMapper(StringifyMapperBase[P]): def map_reduce(self, expr: Any, enclosing_prec: int, *args: P.args, **kwargs: P.kwargs diff --git a/pytato/transform/lower_to_index_lambda.py b/pytato/transform/lower_to_index_lambda.py index 3537b0486..773ee4cbe 100644 --- a/pytato/transform/lower_to_index_lambda.py +++ b/pytato/transform/lower_to_index_lambda.py @@ -36,7 +36,7 @@ import pymbolic.primitives as prim from pymbolic import ArithmeticExpression -from pytools import UniqueNameGenerator +from pytools import UniqueNameGenerator, product from pytato.array import ( AbstractResultWithNamedArrays, @@ -61,7 +61,7 @@ _get_einsum_access_descr_to_axis_len, ) from pytato.diagnostic import CannotBeLoweredToIndexLambda -from pytato.scalar_expr import INT_CLASSES, ScalarExpression +from pytato.scalar_expr import INT_CLASSES, FlattenMapper, ScalarExpression from pytato.tags import AssumeNonNegative from pytato.transform import ( Mapper, @@ -127,13 +127,31 @@ def _generate_index_expressions( new_strides = new_strides[::-1] old_size_tills = old_size_tills[::-1] - flattened_index_expn = sum( + flattened_index_expn = FlattenMapper()(sum( index_var*new_stride - for index_var, new_stride in zip(index_vars, new_strides, strict=True)) + for index_var, new_stride in zip(index_vars, new_strides, strict=True))) + + old_size = product(old_shape) + + def _mod( + num: ArithmeticExpression, denom: ArithmeticExpression + ) -> ArithmeticExpression: + from pymbolic.typing import Integer + if isinstance(old_size, Integer) and denom == old_size and denom != 0: + return num + # Pyright has a point: complex numbers don't support '%'. + return num % denom # pyright: ignore[reportOperatorIssue,reportUnknownVariableType] + + def _floordiv( + num: ArithmeticExpression, denom: ArithmeticExpression + ) -> ArithmeticExpression: + if denom == 1: + return num + # pyright has a point: complex numbers don't support '//'. + return num // denom # pyright: ignore[reportOperatorIssue,reportUnknownVariableType] return tuple( - # Mypy has a point: complex numbers don't support '//'. - (flattened_index_expn % old_size_till) // old_stride # type: ignore[operator] + _floordiv(_mod(flattened_index_expn, old_size_till), old_stride) # pyright: ignore[reportArgumentType] for old_size_till, old_stride in zip(old_size_tills, old_strides, strict=True)) diff --git a/test/test_pytato.py b/test/test_pytato.py index f680c0b27..ca3786e32 100644 --- a/test/test_pytato.py +++ b/test/test_pytato.py @@ -1561,6 +1561,23 @@ def test_lower_to_index_lambda(): assert idx_lambda.expr.index_tuple[4] == Variable("_1") +def test_lower_to_index_lambda_flatten_reshape(): + # Before commit<=577cb39, pytato would generate redundant floor div + # and modulo computations. Specifically, in this case we would get: + # out[_0] = x[((0 + _0*1) % 40) // 4, ((0 + _0*1) % 4) // 1]. + + from pymbolic import parse + + from pytato.array import IndexLambda + x = pt.make_placeholder(name="x", dtype=float, shape=(10, 4)) + idx_lambda = pt.to_index_lambda(x.reshape(-1)) + assert isinstance(idx_lambda, IndexLambda) + assert idx_lambda.expr.index_tuple == ( + parse("_0 // 4"), + parse("_0 % 4"), + ) + + def test_reserved_binding_name_patterns(): from pytato.transform.metadata import BINDING_NAME_RESERVED_PATTERN