jax.extend.core 模块

ClosedJaxpr(jaxpr, consts)
Jaxpr(constvars, invars, outvars, eqns[, ...])
JaxprEqn(invars,  outvars,  primitive,  params,  ...)
Literal(val,  aval)
Primitive(name)
Token(buf)
Var(suffix,  aval)
array_types	set() -> 新的空集合对象 set(iterable) -> 新的集合对象
jaxpr_as_fun
primitives