Extending PyTensor#
This advanced tutorial is for users who want to extend PyTensor with new Types,
new operations (i.e. Ops), and new graph rewrites. This first page of the
tutorial mainly focuses on the Python implementation of an Op and then
proposes an overview of the most important methods that define an Op.
The second page of the tutorial (Extending PyTensor with a C Op) provides then
information on the C implementation of an Op. The rest of the tutorial
goes more in depth on advanced topics related to Ops, such as how to write
efficient code for an Op and how to write an rewrite to speed up the
execution of an Op.
Along the way, this tutorial also introduces many aspects of how PyTensor works, so it is also good for you if you are interested in getting more under the hood with PyTensor itself.
Note
Before tackling this more advanced presentation, it is highly recommended
to read the introductory Tutorial, especially the sections
that introduce the PyTensor graphs, as providing a novel PyTensor Op requires a
basic understanting of the PyTensor graphs.
See also the Developer Start Guide for information regarding the versioning framework, namely about Git and GitHub, regarding the development workflow and how to make a quality contribution.
- Graph Structures
- Graph Rewriting
- Graph and Node Rewriters
- A Simplification Rule
- Unification and reification
- miniKanren
- The Optimization Database (
optdb) - Profiling PyTensor Function Compilation
Ops- Creating a new
Op: Python implementation - Extending PyTensor with a C
Op - Adding JAX, Numba and Pytorch support for
Ops - Overview of the compilation pipeline
Types- Implementing C support for
Types- How does it work?
- What needs to be defined
CLinkerTypeCLinkerType.c_declare()CLinkerType.c_init()CLinkerType.c_extract()CLinkerType.c_sync()CLinkerType.c_cleanup()CLinkerType.c_headers()CLinkerType.c_libraries()CLinkerType.c_header_dirs()CLinkerType.c_lib_dirs()CLinkerType.c_compile_args()CLinkerType.c_no_compile_args()CLinkerType.c_init_code()CLinkerType.c_support_code()CLinkerType.c_compiler()CLinkerType.c_code_cache_version()CLinkerType.c_element_type()
- Defining the methods
- What the generated C will look like
- Final version
DeepCopyOpViewOpShapeandShape_i
- Views and inplace operations
- Developer documentation for
Scan - Implementing some specific Ops
- Using COp params
- Unit Testing
- Extending PyTensor: FAQ and Troubleshooting
- Tips