jax

93c57481 - Tiny usability tweak to allow user to pass a list instead of a tuple, e.g. `pltpu.VMEM([8], jnp.int32)`.