Canonicalize the dtype

Signed-off-by: Ming Huang <[email protected]>
NVIDIA · Nov 3, 2023 · f3a8699 · f3a8699
1 parent 2ab34ee
commit f3a8699
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/transformer_engine/jax/cpp_extensions.py b/transformer_engine/jax/cpp_extensions.py
@@ -84,6 +84,7 @@ def jax_dtype_to_te_dtype(jax_dtype):
     """
     convert jax dtype to TE dtype
     """
+    jax_dtype = dtypes.canonicalize_dtype(jax_dtype)
     if jax_dtype == jnp.float32:
         return TEDType.kFloat32
     if jax_dtype == jnp.float16:
@@ -1071,6 +1072,7 @@ def is_kernel_available(batch: int, heads: int, q_seqlen: int, k_seqlen: int,
         """Check Softmax kernel availability based on size"""
         attn_batches = batch * heads
 
+        dtype = dtypes.canonicalize_dtype(dtype)
         if (dtype in [jnp.float16, jnp.bfloat16]
                 and 16 < k_seqlen <= SoftmaxPrimitive.max_k_seqlen_supported
         # k_seqlen must be 16 ~ 4096
@@ -1225,6 +1227,7 @@ def is_kernel_available(batch: int, heads: int, q_seqlen: int, k_seqlen: int,
         """Check Softmax kernel availability based on size"""
         attn_batches = batch * heads
 
+        dtype = dtypes.canonicalize_dtype(dtype)
         if (dtype in [jnp.float16, jnp.bfloat16]
                 and 16 < k_seqlen <= SoftmaxPrimitive.max_k_seqlen_supported
         # k_seqlen must be 16 ~ 4096
@@ -1441,6 +1444,7 @@ def is_kernel_available(batch: int, heads: int, q_seqlen: int, k_seqlen: int,
         """Check Softmax kernel availability based on size"""
         attn_batches = batch * heads
 
+        dtype = dtypes.canonicalize_dtype(dtype)
         if (dtype in [jnp.float16, jnp.bfloat16]
                 and 16 < k_seqlen <= SoftmaxPrimitive.max_k_seqlen_supported
         # k_seqlen must be 16 ~ 4096