jax

b1981cbb - [hijax] improved dce interface

[hijax] improved dce interface * pytrees used_inputs and used_outputs * the rule can decide what outputs it's going to produce * don't expose jaxpr eqn, instead just use HiPrimitive instances Co-authored-by: Yash Katariya <yashkatariya@google.com>