Implement all Ops in PyTorch (help welcome!)

Description

If you want to help implementing some of these Ops just leave a comment below saying which ones you are interested in. We'll give you some time to work on it (and then put it back up to grabs).

See the documentation for How to implement PyTorch Ops and tests: https://pytensor.readthedocs.io/en/latest/extending/creating_a_numba_jax_op.html

Example PR: #836

See https://github.com/pymc-devs/pytensor/issues/821#issuecomment-2202258929 for suggestions on equivalent torch functions

Tensor creation Ops

• Alloc and AllocEmpty #836
• Arange #836
• Eye #877
• ScalarFromTensor
• TensorFromScalar
• Repeat #890
• Unique #890
• Sort / Argsort #897
• Tri

Shape Ops

• Dimshuffle - Done in #764
• Reshape #926
• Shape, Shape_i #926
• SpecifyShape #926
• Unbroadcast #926
• Join #869
• Split
• MakeVector #926

Math Ops

• Elemwise - Done in #764 (but not complete! Specific cases will require a custom ScalarOp dispatch)
• CAReduce (Sum, All, Any...) - Assigned to @HarshvirSandhu
• CumOp - #837
• Softmax, LogSoftmax and Grads in #846
• Dot #878
• BatchedDot #878
• Argmax

Indexing Ops

• Subtensor #910
• Inc/SetSubtensor #910
• AdvancedSubtensor[1] #910
• AdvancedIncSubtensor[1] #910

Branching Ops

• CheckAndRaise - #764
• Ifelse #940
• ScalarLoop #958
• Scan
• OpFromGraph #956
• Blockwise

Linalg Ops

• SVD #920
• Det #920
• Eig #920
• Eigh #920
• MatrixInverse #920
• MatrixPinv #920
• QRFull #920
• SLogDet #920
• BlockDiagonal #922
• Cholesky #922
• Solve #922
• SolveTriangular #922

SparseOps

• ... (to be filled)

RandomVariable Ops

• ... Need to figure out API differences

If you need an Op that's not in this list, comment below and we'll add it!

Choose one unchecked operation from the issue, then read the linked guide on implementing PyTorch Ops and inspect example PR #836. Use the equivalent Torch functions noted in the issue comments; done means the selected operation is implemented with tests and its checklist entry is resolved.