Extending Aesara

This advanced tutorial is for users who want to extend Aesara with new Types, new Operations (Ops), and new graph optimizations. 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 Aesara 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 optimization to speed up the execution of an Op.

Along the way, this tutorial also introduces many aspects of how Aesara works, so it is also good for you if you are interested in getting more under the hood with Aesara itself.

Note

Before tackling this more advanced presentation, it is highly recommended to read the introductory Tutorial, especially the sections that introduce the Aesara Graphs, as providing a novel Aesara op requires a basic understanting of the Aesara 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.

• Creating a new Op: Python implementation
  • Aesara Graphs refresher
  • Op’s basic methods
  • Op’s auxiliary methods
  • Example: Op definition
  • Example: __props__ definition
  • How To Test it
    • Basic Tests
    • Testing the infer_shape
    • Testing the gradient
    • Testing the Rop
    • Running Your Tests
      • In-file
      • Exercise
      • Random numbers in tests
  • as_op()
    • as_op() Example
    • Exercise
  • Documentation and Coding Style
  • NanGuardMode and AllocEmpty
  • Final Note
• Extending Aesara with a C Op
  • Python C-API
    • Reference counting
  • NumPy C-API
    • NumPy data types
    • NumPy ndarrays
    • Accessing NumPy ndarrays’ data and properties
    • Creating NumPy ndarrays
  • Methods the C Op needs to define
  • Important restrictions when implementing a COp
  • Simple COp example
  • More complex C Op example
  • Alternate way of defining C Ops
    • Main function
    • Macros
    • Support code
  • Using GDB to debug COp’s C code
  • Final Note
• Writing an Op to work on an ndarray in C
  • Writing an Optimization
  • Test the optimization
• Overview of the compilation pipeline
  • Definition of the computation graph
  • Compilation of the computation graph
    • Step 1 - Create a FunctionGraph
    • Step 2 - Execute main Optimizer
    • Step 3 - Execute linker to obtain a thunk
    • Step 4 - Wrap the thunk in a pretty package
• Aesara vs. C
• Graph Structures
  • Traversing the graph
  • Graph Structures
    • Apply
    • Op
    • Type
    • Variable
      • Constant
  • Graph Structures Extension
  • Automatic Differentiation
  • Optimizations
• Making the double Type
  • Type’s contract
    • Additional definitions
      • transfer()
  • Defining double
  • Untangling some concepts
  • Final version
• Making arithmetic Ops on double
  • Op’s contract
  • Optional methods or attributes
  • Gradient
  • Defining an Op: mul
  • Trying out our new Op
    • Automatic Constant Wrapping
  • Final version
• Views and inplace operations
  • Views
  • Inplace operations
  • Destructive Operations
  • Inplace optimization and DebugMode
• Implementing some specific Ops
  • Scalar/Elemwise/Reduction Ops
  • Sparse Ops
    • Sparse Gradient
    • Sparse C code
    • Sparse Tests
  • Random distribution
  • OpenMP Ops
  • Numba Ops
  • Alternate Aesara Types
• Implementing double in C
  • How does it work?
  • What needs to be defined
  • Defining the methods
  • What the generated C will look like
  • Final version
  • DeepCopyOp
  • ViewOp
  • Shape and Shape_i
• Implementing the arithmetic COps in C
  • How does it work?
  • What needs to be defined
  • Defining the methods
  • Final version
• Using Op params
  • The params type
  • The params object
    • Defining a params type
  • Registering the params with your Op
  • Signature changes from having params
  • A simple example
  • A more complex example
• Graph optimization
  • Global and Local Optimizations
    • Global optimization
    • Local optimization
  • One simplification rule
    • Global optimization
    • Local Optimization
      • OpSub, OpRemove, PatternSub
  • The optimization database (optdb)
    • Definition of optdb
    • OptimizationQuery
    • Examples
    • Registering an Optimizer
    • Registering a LocalOptimizer
    • optdb structure
    • NavigatorOptimizer
  • Profiling Aesara function compilation
    • Detailed profiling of Aesara optimizations
• Tips
  • Reusing outputs
  • Don’t define new Ops unless you have to
  • Use Aesara’s high order Ops when applicable
  • Op Checklist
• Unit Testing
  • How to Run Unit Tests
  • Folder Layout
    • How to Write a Unit Test
  • Test Cases and Methods
  • Checking for correctness
    • Creating an Op Unit Test
  • Validating the Gradient
    • makeTester and makeBroadcastTester
• Developer documentation for Scan
  • Context
  • Pre-requisites
  • Relevant code files
  • Notation
  • Scan variables
  • Optimizations
    • remove_constants_and_unused_inputs_scan
    • PushOutNonSeqScan
    • PushOutSeqScan
    • PushOutScanOutput
    • PushOutDot1
    • ScanInplaceOptimizer
    • ScanSaveMem
    • ScanMerge
    • scan_merge_inouts
  • Helper classes and functions
    • Accessing/manipulating Scan’s inputs and outputs by type
    • Navigating between outer inputs/outputs and inner inputs/outputs
• Extending Aesara: FAQ and Troubleshooting
  • I wrote a new Op/Type, and weird stuff is happening…
  • I wrote a new optimization, but it’s not getting used…
  • I wrote a new optimization, and it changed my results even though I’m pretty sure it is correct.
• Tutorial on adding JAX Ops to Aesara
• Step 1: Identify the Aesara Op you’d like to JAXify
• Step 2: Find the relevant JAX method (or something close)
• Step 3: Register the function with the jax_funcify dispatcher
• Step 4: Write tests
• Step 5: Wait for CI pass and Code Review
• Appendix: What does singledispatcher do?