Extending Aesara

This advanced tutorial is for users who want to extend Aesara 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 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 rewrite 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.

• Graph Structures
  • Traversing the graph
  • Graph Structures
    • Apply
    • Op
    • Type
    • Variable
      • Constant
  • Automatic Differentiation
  • Rewrites
• Graph Rewriting
  • Graph and Node Rewriters
    • Graph Rewriting
      • GraphRewriter
        • GraphRewriter.apply()
        • GraphRewriter.add_requirements()
        • GraphRewriter.rewrite()
    • Node Rewriting
      • NodeRewriter
        • NodeRewriter.transform()
  • A Simplification Rule
    • Graph Rewriter Implementation
    • Node Rewriter Implementation
      • SubstitutionNodeRewriter, RemovalNodeRewriter, PatternNodeRewriter
        • SubstitutionNodeRewriter()
        • RemovalNodeRewriter()
        • PatternNodeRewriter()
  • Unification and reification
  • miniKanren
  • The Optimization Database (optdb)
    • Definition of optdb
    • RewriteDatabaseQuery
      • RewriteDatabaseQuery
        • RewriteDatabaseQuery.include
        • RewriteDatabaseQuery.require
        • RewriteDatabaseQuery.exclude
        • RewriteDatabaseQuery.subquery
    • Examples
    • Registering a Rewriter
    • Registering a NodeRewriter
      • canonicalize()
      • specialize()
    • optdb structure
    • NodeProcessingGraphRewriter
  • Profiling Aesara Function Compilation
    • Detailed Profiling of Aesara Rewrites
• Ops
  • make_node()
  • perform()
  • __eq__()
  • __hash__()
  • Optional methods or attributes
    • __props__
    • default_output
    • make_thunk()
    • __call__()
    • infer_shape()
    • flops()
    • __str__()
    • do_constant_folding()
    • debug_perform()
  • Gradient
    • grad()
    • R_op()
• 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 Op.infer_shape()
    • Testing the gradient
    • Testing the Rop
    • Running Your Tests
      • In-file
      • Exercise
      • Random numbers in tests
  • as_op()
    • as_op() Example
  • 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 :class`ndarray`s’ data and properties
    • Creating NumPy ndarrays
  • Methods the C Op needs to define
    • c_code()
    • c_support_code()
    • c_support_code_apply()
    • c_code_cache_version()
  • 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
• Adding JAX and Numba support for Ops
  • Step 1: Identify the Aesara Op you’d like to implement in JAX
  • Step 2: Find the relevant JAX method (or something close)
  • Step 3: Register the function with the _jax_funcify dispatcher
  • Step 4: Write tests
• Overview of the compilation pipeline
  • Step 1 - Create a FunctionGraph
  • Step 2 - Perform graph rewrites
  • Step 3 - Execute linker to obtain a thunk
  • Step 4 - Wrap the thunk in a pretty package
• Types
  • The TensorType
    • Type’s contract
  • Additional definitions
    • An example
    • Untangling some concepts
• Implementing C support for Types
  • 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 and Shape_i
• Views and inplace operations
  • Views
  • Inplace operations
  • Destructive Operations
  • Inplace Rewriting and DebugMode
• Developer documentation for Scan
  • Context
  • Pre-requisites
  • Relevant code files
  • Notation
  • Scan variables
  • Rewrites
    • 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
• 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
• 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
• 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()
• Extending Aesara: FAQ and Troubleshooting
  • I wrote a new Op/Type, and weird stuff is happening…
  • I wrote a new rewrite, but it’s not getting used…
  • I wrote a new rewrite, and it changed my results even though I’m pretty sure it is correct.
• Tips
  • Don’t define new Ops unless you have to
  • Use Aesara’s high order Ops when applicable