jax
566d0775 - unify `stages.Lowering` and `stages.XlaLowering`

Commit
1 year ago
unify `stages.Lowering` and `stages.XlaLowering` We no longer have many different implicit types conforming to `Lowering`, only `pxla.MeshComputation` and `pxla.PmapComputation`. Both are `XlaLowering` subtypes. So define just one common base class, call it `Lowering`, and inherit from just that in both concrete internal computation/lowering subtypes. PiperOrigin-RevId: 746735857
Author
Parents
Loading