[JAX] Grouped GEMM Refactor to use first_dims and last_dims

tests/jax/test_custom_call_compute.py

@@ -1787,13 +1787,19 @@ def test_grouped_gemm_fp16(self, dtype, input_shape, layout):
         ref_out = self._ref_grouped_dense(lhs, rhs, None, group_sizes, contracting_dims)
 
         # jitting grouped_gemm
+        empty_gs = jnp.empty((0,), jnp.int32)
         prim_out = jax.jit(
             tex.grouped_gemm, static_argnames=("contracting_dims", "use_async_d2h_group_sizes")
         )(
             lhs,
             rhs,
-            group_sizes,
-            contracting_dims,
+            lhs_first_dims=group_sizes,
+            lhs_last_dims=empty_gs,
+            rhs_first_dims=empty_gs,
+            rhs_last_dims=empty_gs,
+            out_first_dims=group_sizes,
+            out_last_dims=empty_gs,
+            contracting_dims=contracting_dims,
             use_async_d2h_group_sizes=True,
         )
 
@@ -1825,8 +1831,18 @@ def test_grouped_gemm_fp8(self, fwd_bwd_dtype, scaling_mode, input_shape, layout
         )
         ref_out = self._ref_grouped_dense(lhs, rhs, None, group_sizes, contracting_dims)
 
+        empty_gs = jnp.empty((0,), jnp.int32)
         prim_out = jax.jit(tex.grouped_gemm, static_argnames=("contracting_dims",))(
-            lhs, rhs, group_sizes, contracting_dims, quantizer_set=quantizer_set
+            lhs,
+            rhs,
+            lhs_first_dims=group_sizes,
+            lhs_last_dims=empty_gs,
+            rhs_first_dims=empty_gs,
+            rhs_last_dims=empty_gs,
+            out_first_dims=group_sizes,
+            out_last_dims=empty_gs,
+            contracting_dims=contracting_dims,
+            quantizer_set=quantizer_set,
         )
 
         allclose_dtype = jnp.float8_e4m3fn