Skip to content

Conversation

@tchan102
Copy link
Contributor

@tchan102 tchan102 commented Nov 4, 2025

Add optimization for Join → Repeat when concatenating identical tensors

Description

This PR introduces a graph rewrite optimization in pytensor/tensor/rewriting/basic.py that replaces redundant Join operations with an equivalent and more efficient Repeat operation when all concatenated tensors are identical.

Example:
join(0, x, x, x) → repeat(x, 3, axis=0)

Key additions:

  • Implemented new rewrite function local_join_to_repeat registered under both @register_canonicalize and @register_specialize.
  • Added corresponding test test_local_join_to_repeat to verify correctness, performance, and behavior for vectors and matrices.

Related Issue

Checklist

Type of change

  • [ x] New feature / enhancement
  • Bug fix
  • Documentation
  • Maintenance
  • Other (please specify):

@ricardoV94 ricardoV94 added graph rewriting enhancement New feature or request labels Nov 4, 2025
@ricardoV94
Copy link
Member

Let's try with @register_canonicalize only

@ricardoV94
Copy link
Member

Btw would be nice to get rid of this join (and split) symbolic axis if you would like to work on that after this PR. relevant issue: #1528

@tchan102 tchan102 requested a review from ricardoV94 November 9, 2025 00:47
@tchan102 tchan102 requested a review from ricardoV94 November 11, 2025 20:51
@tchan102 tchan102 requested a review from ricardoV94 November 12, 2025 07:46
@tchan102 tchan102 requested a review from ricardoV94 November 13, 2025 01:40
@ricardoV94 ricardoV94 merged commit 0b731c2 into pymc-devs:main Nov 14, 2025
53 of 56 checks passed
@ricardoV94 ricardoV94 changed the title Rewrite concatenate([x, x]) as repeat(x, 2) Rewrite concatenate([x, x]) as tile Nov 14, 2025
@ricardoV94
Copy link
Member

Thanks @tchan102

Took a bit longer than expected because it touched some previous tests

@tchan102
Copy link
Contributor Author

Thanks @ricardoV94 for reviewing this pull request and checking my errors, it's my first time contributing an open source project. I'm really appreciate your help and guidance. Looking forward to push more code to PyTensor!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancement New feature or request graph rewriting

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Rewrite concatenate([x, x]) as repeat(x, 2)

2 participants