Add JAX unit test for Shardy which causes the compiler to introduce the mlir::tensor::TensorDialect. This was causing the compiler to crash.

Add JAX unit test for Shardy which causes the compiler to introduce the mlir::tensor::TensorDialect. This was causing the compiler to crash. #2856