API Reference

This reference gives the public signatures, parameters, and return values for tracing, transformation, execution, model specifications, and extension interfaces. The pages group related APIs for lookup. A First Program introduces these interfaces, while the concept pages explain the behavior behind the signatures.

• Core
  • trace()
• IR Transforms
  • batch()
  • pullback()
  • pushforward()
  • sched()
  • dce()
  • weight()
• Context Managers
  • fold()
  • tag()
  • memoize()
  • collect()
  • inject()
  • client()
  • LM Clients
    • Client
    • LiteLLMClient
• Custom Boundaries
  • custom()
• Primitives
  • fill()
  • String
    • format()
    • concat()
    • match()
  • Numeric
    • neg()
    • add()
    • sub()
    • mul()
    • div()
    • eq()
    • ne()
    • lt()
    • le()
    • gt()
    • ge()
  • Control Flow
    • stop_gradient()
    • switch()
    • while_loop()
    • fixpoint()
  • Scheduling
    • depends()
  • Intercepts
    • checkpoint()
  • Trace Weight
    • factor()
• Language Models
  • Specifications
  • Structured Output
  • Model Clients
• Schemas
  • Str
  • Int
  • Float
  • Bool
  • Enum
  • Schema Utilities
    • describe()
    • parse()
• Types
  • autoform.__version__
  • autoform.PYTREE_NAMESPACE
• Extend
  • Types
    • AVal
    • StrAVal
    • IntAVal
    • FloatAVal
    • BoolAVal
    • Primitives
    • Dunders
    • IR
    • Interpreters
  • Rule Registration
    • register_impl()
    • register_aimpl()
    • register_abstract()
    • register_batch()
    • register_abatch()
    • register_pushforward()
    • register_apushforward()
    • register_pullback_fwd()
    • register_apullback_fwd()
    • register_pullback_bwd()
    • register_apullback_bwd()
    • register_dce()
  • Abstract Helpers
    • Space
    • primal_s
    • tangent_s
    • cotangent_s
    • materialize_zeros()
  • Batch Helpers
    • batch_index()
    • batch_spec()
    • batch_transpose()
  • IR Helpers
    • using_interpreter()
    • serial_fanout()
    • is_var()
    • aval_if_var()
  • Active Contexts
    • active_interpreter
    • active_tags
    • active_client
  • Primitive Keys
    • concat_p
    • match_p
    • neg_p
    • add_p
    • sub_p
    • mul_p
    • div_p
    • eq_p
    • ne_p
    • lt_p
    • le_p
    • gt_p
    • ge_p
    • fill_p
    • checkpoint_p
    • stop_gradient_p
    • switch_p
    • while_loop_p
    • fixpoint_p
    • factor_p
    • weight_call_p
    • fanout_p
    • depends_p
    • batch_call_p
    • pushforward_call_p
    • pullback_call_p
  • Registration
    • register_trace_type()
    • register_non_dce()
    • register_non_memoizable()
  • Dunder Registration
    • register_dunder()