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]

one thing you'll need to handle is that join axis is symbolic but repeat must be a python integer. You should do something like:

if not isinstance(axis, Constant):
    return None  # rewrite only applies to constant axis
axis = axis.data

[ricardoV94]

The RepeatOp will not actually be used. It's only used when there are vector repeats. For scalar Repeats we end up with Alloc instead (PyTensor version of BroadcastTo). Did you try to run the tests locally?

[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