Skip to content

Commit cfee66b

Browse files
committed
fix test
1 parent 5749e4c commit cfee66b

File tree

1 file changed

+6
-0
lines changed

1 file changed

+6
-0
lines changed

torch_xla/experimental/custom_kernel.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -183,6 +183,12 @@ def convert_torch_dtype_to_jax(dtype: torch.dtype) -> "jnp.dtype":
183183
return jnp.int8
184184
elif dtype == torch.uint8:
185185
return jnp.uint8
186+
elif dtype == torch.float8_e5m2:
187+
return jnp.float8_e5m2
188+
elif dtype == torch.float8_e4m3fn:
189+
return jnp.float8_e4m3fn
190+
elif dtype == torch.float8_e4m3fnuz:
191+
return jnp.float8_e4m3fnuz
186192
else:
187193
raise ValueError(f"Unsupported dtype: {dtype}")
188194

0 commit comments

Comments
 (0)