Add tests

[nstarman]

• unit tests
• test whether something can be more directly imported from jax.experimental.array_api. In particular, this repo adds jaxtyping + beartype runtime-checking support, which relies on type hints that are not in jax.experimental.array_api.