New transform: Push indexing to only materialized nodes.

[kaushikcfd]

Generalizes the logic in #425.

Includes commits from #656.

[Copilot]

Pull request overview

Adds a new transform to push indexing operations down the expression DAG so that indexing happens only at (or is forced onto) materialized nodes, generalizing prior work in #425. This is accompanied by a substantial new test suite and a small testing helper to compare transformed vs. reference evaluation.

Changes:

• Introduce pt.push_index_to_materialized_nodes transform and export it from pytato.__init__.
• Add extensive tests covering many indexing composition/broadcasting scenarios.
• Add assert_allclose_to_ref test helper for comparing two Pytato expressions by evaluation.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 4 comments.

Show a summary per file

File	Description
pytato/transform/push_index_to_materialized_nodes.py	New transform implementation that composes/pushes indices through various ops.
pytato/__init__.py	Exposes push_index_to_materialized_nodes as part of the public API.
pytato/raising.py	Adjusts ZerosLikeOp typing/raising for the new transform’s HLO lowering path.
pytato/array.py	Tweaks index-expression typing (ConvertibleToIndexExpr).
test/testlib.py	Adds assert_allclose_to_ref helper to compare two Pytato expressions via execution.
test/test_transform.py	New large test module for the indexing pusher transform.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.