add lazy JAX arraycontext (#388)

* add lazy JAX arraycontext

Co-authored-by: Alex Fikl <alexfikl@gmail.com>

* typing fixes

* workaround bpr

* bpr fixes

* modify nitpick ignores

---------

Co-authored-by: Alex Fikl <alexfikl@gmail.com>