brainstate.transform.eqns_to_jaxpr

Contents

brainstate.transform.eqns_to_jaxpr#

brainstate.transform.eqns_to_jaxpr(eqns, invars=None, outvars=None, constvars=None)[source]#

Convert a sequence of JaxprEqn into a Jaxpr.

Parameters:

  • eqns (Sequence[JaxprEqn]) – Sequence of Jaxpr equations to convert

  • invars (Sequence[Var]) – Input variables. If None, will be inferred from equations

  • outvars (Sequence[Var]) – Output variables. If None, will be inferred from equations

  • constvars (Sequence[Var]) – Constant variables. If None, will be automatically extracted from equations

Returns:

A Jaxpr object constructed from the equations

Return type:

Jaxpr

Contents