Implement all Ops in PyTorch (help welcome!)

[ricardoV94]

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 #821 (comment) for suggestions on equivalent torch functions

Tensor creation Ops

• Alloc and AllocEmpty Implement basic Alloc Ops in PyTorch #836
• Arange Implement basic Alloc Ops in PyTorch #836
• Eye Implemented Eye Op in PyTorch #877
• ScalarFromTensor
• TensorFromScalar
• Repeat Implemented Repeat and Unique Ops in PyTorch #890
• Unique Implemented Repeat and Unique Ops in PyTorch #890
• Sort / Argsort Implemented Sort/Argsort Ops in PyTorch #897
• Tri

Shape Ops

• Dimshuffle - Done in Add initial support for PyTorch backend #764
• Reshape Implements shape Ops and MakeVector in PyTorch #926
• Shape, Shape_i Implements shape Ops and MakeVector in PyTorch #926
• SpecifyShape Implements shape Ops and MakeVector in PyTorch #926
• Unbroadcast Implements shape Ops and MakeVector in PyTorch #926
• Join Pytorch support for Join and Careduce Ops #869
• Split
• MakeVector Implements shape Ops and MakeVector in PyTorch #926

Math Ops

• Elemwise - Done in Add initial support for PyTorch backend #764 (but not complete! Specific cases will require a custom ScalarOp dispatch)
• CAReduce (Sum, All, Any...) - Assigned to @HarshvirSandhu
• CumOp - Add docs on implementing Pytorch Ops (and CumOp) #837
• Softmax, LogSoftmax and Grads in PyTorch Softmax Ops #846
• Dot Implement Dot and BatchedDot in PyTensor #878
• BatchedDot Implement Dot and BatchedDot in PyTensor #878
• Argmax

Indexing Ops

• Subtensor Implement indexing operations in PyTorch #910
• Inc/SetSubtensor Implement indexing operations in PyTorch #910
• AdvancedSubtensor[1] Implement indexing operations in PyTorch #910
• AdvancedIncSubtensor[1] Implement indexing operations in PyTorch #910

Branching Ops

• CheckAndRaise - Add initial support for PyTorch backend #764
• Ifelse Add torch conditionals: IfElse #940
• ScalarLoop Implement ScalarLoop in torch backend #958
• Scan
• OpFromGraph Implement OpFromGraph in PyTorch backend #956
• Blockwise

Linalg Ops

• SVD Implement nlinalg Ops in PyTorch #920
• Det Implement nlinalg Ops in PyTorch #920
• Eig Implement nlinalg Ops in PyTorch #920
• Eigh Implement nlinalg Ops in PyTorch #920
• MatrixInverse Implement nlinalg Ops in PyTorch #920
• MatrixPinv Implement nlinalg Ops in PyTorch #920
• QRFull Implement nlinalg Ops in PyTorch #920
• SLogDet Implement nlinalg Ops in PyTorch #920
• BlockDiagonal Implement linear algebra functions in PyTorch #922
• Cholesky Implement linear algebra functions in PyTorch #922
• Solve Implement linear algebra functions in PyTorch #922
• SolveTriangular Implement linear algebra functions in PyTorch #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!

[HAKSOAT]

Hi there @ricardoV94, I'm attending the hackathon at Pydata London.

I'd like to work on Softmax.

[ricardoV94]

Hi there @ricardoV94, I'm attending the hackathon at Pydata London.

I'd like to work on Softmax.

Sure!

[HAKSOAT]

Thanks @ricardoV94

I have been doing some reading around the docs, I noticed that #764 has the initial setup. I believe that contains ground work for adding other Ops, so I'd be following the PR and the Softmax Ops will likely use some code from it.

Am I thinking about this the right way?

[ricardoV94]

@HAKSOAT yes, you should be able to start a branch already from #764. All the functionality is done, except some tests are failing because they depend non-optionally on GPU. We should make that optional and get it merged soon enough, but you need not wait :)

[ricardoV94]

#764 is merged

[t3chw]

Hello @ricardoV94, I would like to work on Reshape

[ricardoV94]

Go ahead. We'll link and lock the Op when you open a PR

[HAKSOAT]

Hi @ricardoV94, I have opened a PR for the Softmax Ops. I see it has been grouped with LogSoftmax and Grads, so I can update the PR to include them.