API ReferenceΒΆ Core trace() IR Transforms batch() pullback() pushforward() sched() dce() weighted() Context Managers fold() tag() memoize() collect() inject() lm_client() Custom Boundaries custom() Primitives LM lm_call() lm_schema_call() String format() concat() match() Control Flow stop_gradient() switch() while_loop() Scheduling depends() Intercepts checkpoint() Trace Weight factor() Schemas Str Int Float Bool Enum Doc Types autoform.__version__ autoform.PYTREE_NAMESPACE Extend 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()