jax

7ff389bd - extend type transfer to all primitives, including call and map primitives