JAX Matrix

[kels271828]

For JAX models to play nicely with other packages that use msca.linalg.matrix.Matrix subclasses, we might want to consider creating a msca.linalg.matrix.JAXMatrix.

[zhengp0]

Thanks @kels271828! I think people have been thinking about this
jax-ml/jax#16290
There might be difficulties for directly inherit the jax.Array unfortunately.