Update jax gemm.py

Author: alextmagro

transformer_engine/jax/cpp_extensions/gemm.py

@@ -20,19 +20,22 @@
 from jax.sharding import NamedSharding, PartitionSpec
 from jax.experimental.custom_partitioning import SdyShardingRule
 
+from ..util import is_hip_extension, get_jnp_float8_e4m3_type, get_jnp_float8_e5m2_type
+
 from transformer_engine_jax import (
     get_num_compute_streams,
     JAXX_Collective_Op,
     get_device_compute_capability,
-    #initialize_cgemm_communicator,
-    #get_cgemm_num_max_streams,
 )
+if not is_hip_extension():
+    from transformer_engine_jax import (
+        initialize_cgemm_communicator,
+        get_cgemm_num_max_streams,
+    )
 
 from .base import BasePrimitive, register_primitive
 from .quantization import grouped_quantize
 
-from ..util import is_hip_extension, get_jnp_float8_e4m3_type, get_jnp_float8_e5m2_type
-
 from ..quantize import (
     AbstractBaseTensor,
     NoScaleTensor,