mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-16 03:46:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

[Pallas][Mosaic] Support float8_e4m3b11fnuz

Browse Source 
PiperOrigin-RevId: 729169181
This commit is contained in:
jax authors 2025-02-20 10:43:46 -08:00
parent ddcb7deeaf
commit b7968474c2
2 changed files with 4 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
jax/_src/pallas/mosaic/lowering.py
View File

@ -1851,6 +1851,7 @@ def _dot_general_lowering_rule(
ir.F32Type,
ir.Float8E5M2Type,
ir.Float8E4M3FNType,
ir.Float8E4M3B11FNUZType,
]
):
val = ir.FloatAttr.get(val_type, 0.0)

8
jaxlib/mosaic/dialect/tpu/util.h
View File

@ -162,13 +162,11 @@ FailureOr<int8_t> getTypeBitwidth(Type ty) {
if (auto bf16_ty = dyn_cast<BFloat16Type>(ty)) {
return 16;
}
if (auto f8e5m2_ty = dyn_cast<Float8E5M2Type>(ty)) {
if (isa<Float8E5M2Type, Float8E4M3FNType, Float8E4M3B11FNUZType>(ty)) {
return 8;
}
if (auto f8e4m3fn_ty = dyn_cast<Float8E4M3FNType>(ty)) {
return 8;
}
return emitError(UnknownLoc::get(ty.getContext()), "Unsupported type: ")
return emitError(UnknownLoc::get(ty.getContext()),
"Unsupported type in mosaic dialect: ")
<< ty;
}