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

[tchan102]

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

• [ x] Closes Rewrite concatenate([x, x]) as repeat(x, 2) #1710

Checklist

• [x ] Checked that the pre-commit linting/style checks pass
• [x ] Included tests that prove the fix is effective or that the new feature works
• [x ] Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

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

[ricardoV94]

Let's try with @register_canonicalize only

[ricardoV94]

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

[ricardoV94]

Why? I don't think this is needed condition

[tchan102]

Thanks for the feedback! I wanted to clarify about removing the shape[axis_val] != 1 condition.

When I removed this check, the tests fail because join and repeat behave differently when the size along the join axis is not 1:

Example:
x = vector("x") # e.g., [1.0, 2.0]

Case 1: Without ExpandDims (size is NOT 1 along axis 0)
join(0, x, x) # → [1.0, 2.0, 1.0, 2.0] (concatenates)
repeat(x, 2, 0) # → [1.0, 1.0, 2.0, 2.0] (repeats each element)
❌ These are NOT equivalent!

Case 2: With ExpandDims (size IS 1 along axis 0)
join(0, x[None], x[None]) # → [[1.0, 2.0], [1.0, 2.0]]
repeat(x[None], 2, 0) # → [[1.0, 2.0], [1.0, 2.0]]
✅ These ARE equivalent!

The optimization seems to only be mathematically correct when the size along the join axis is 1 (e.g., after ExpandDims). Without this check, test_local_join_to_repeat fails.

Could you clarify what you meant by "not needed"? Should we:

1. Keep the check as-is (only optimize when shape[axis_val] == 1)
2. Remove only the defensive parts (shp is None or axis_val >= len(shp))
3. Apply the optimization more broadly (and update the tests accordingly)

Thanks for your guidance!