ExtendΒΆ The extension API exposes the low-level surface for adding value domains, primitive behavior, and transform rules. Types Abstract Values Primitives IR Interpreters Rule Registries impl_rules abstract_rules push_rules pull_fwd_rules pull_bwd_rules batch_rules Abstract Helpers avalof() zeroof() materialize() is_zero() Batch Helpers batch_index() batch_spec() batch_transpose() IR Helpers using_interpreter() serial_fanout() is_irvar() ir_aval() Active Contexts active_interpreter active_tags active_client Primitive Keys format_p concat_p match_p lm_call_p lm_schema_call_p checkpoint_p stop_gradient_p switch_p while_loop_p gather_p depends_p batch_call_p pushforward_call_p pullback_call_p Registration register_trace_type() register_zero() register_cotangent_accumulator() register_non_dce() register_non_memoizable() Operator Registration register_add() register_sub() register_mul() register_div() register_matmul() register_eq()