Extending PyTensor#
This advanced tutorial is for users who want to extend PyTensor with new Type
s,
new operations (i.e. Op
s), 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 Op
s, 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
Op
s- Creating a new
Op
: Python implementation - Extending PyTensor with a C
Op
- Adding JAX and Numba support for
Op
s - Overview of the compilation pipeline
Type
s- Implementing C support for
Type
s- How does it work?
- What needs to be defined
CLinkerType
CLinkerType.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
DeepCopyOp
ViewOp
Shape
andShape_i
- Views and inplace operations
- Developer documentation for
Scan
- Implementing some specific Ops
- Using COp params
- Unit Testing
- Extending PyTensor: FAQ and Troubleshooting
- Tips