[JAX] Set `precision=HIGHEST` for the ref_grouped_gemm impl in the unit test (#1967)

* set precision=HIGHEST for the ref_grouped_gemm impl in the unit test
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>


---------
Signed-off-by: Phuong Nguyen <phuonguyen@nvidia.com>

tests/jax/test_custom_call_compute.py

@@ -1265,7 +1265,9 @@ class TestGroupedDense:
         ref_out = []
         dim_num = (contracting_dims, ((), ()))
         for lhs_i, rhs_i, bias_i in zip(lhs, rhs, bias):
-            out_i = jax.lax.dot_general(lhs_i, rhs_i, dim_num) + jnp.expand_dims(bias_i, axis=0)
+            out_i = jax.lax.dot_general(
+                lhs_i, rhs_i, dim_num, precision=jax.lax.Precision.HIGHEST
+            ) + jnp.expand_dims(bias_i, axis=0)
             ref_out.append(jnp.squeeze(out_i))
         return ref_out