[JAX] Add Shardy warning in GEMM custom call (#2101)

* added shardy warning

Signed-off-by: Phuong Nguyen <[email protected]>


---------

Signed-off-by: Phuong Nguyen <[email protected]>

Author: phu0ngng

transformer_engine/jax/cpp_extensions/gemm.py

@@ -8,6 +8,7 @@
 from collections.abc import Iterable
 from typing import Tuple, Sequence, Union
 from functools import partial, reduce
+import warnings
 
 import jax
 import jax.numpy as jnp
@@ -658,6 +659,12 @@ def shardy_sharding_rule(
 
         prefix = "GemmPrimitive_"
 
+        warnings.warn(
+            "Known issues with TE GemmPrimitives when Shardy propagation is enabled. For now,"
+            " please turn off Shardy by exporting the environment variable"
+            " 'JAX_USE_SHARDY_PARTITIONER=0' if you experience any problems."
+        )
+
         def _generate_operand_rules(name, ndim, cdims):
             specs = []
             ldims = tuple(i for i in range(ndim) if i not in cdims)