Add numerics config for FP5 with subnormals

PiperOrigin-RevId: 671841596

Author: ychzhang

aqt/jax/v2/numerics/fp_numerics.py

@@ -31,6 +31,9 @@ class FpNumericsConfig:
   has_two_nan: bool
   has_naninf: bool
 
+##########
+## FP6
+##########
 e5m0_ocp = FpNumericsConfig(
     nexp=5,
     minexp=0,
@@ -85,6 +88,57 @@ class FpNumericsConfig:
     has_naninf=False,
 )
 
+##########
+## FP5
+##########
+e4m0_ocp = FpNumericsConfig(
+    nexp=4,
+    minexp=0,
+    nmant=0,
+    has_subnormals=True,
+    has_two_nan=False,
+    has_naninf=False,
+)
+
+e3m1_ocp = FpNumericsConfig(
+    nexp=3,
+    minexp=0,
+    nmant=1,
+    has_subnormals=True,
+    has_two_nan=False,
+    has_naninf=False,
+)
+
+e2m2_ocp = FpNumericsConfig(
+    nexp=2,
+    minexp=0,
+    nmant=2,
+    has_subnormals=True,
+    has_two_nan=False,
+    has_naninf=False,
+)
+
+e1m3_ocp = FpNumericsConfig(
+    nexp=1,
+    minexp=0,
+    nmant=3,
+    has_subnormals=True,
+    has_two_nan=False,
+    has_naninf=False,
+)
+
+e0m4_ocp = FpNumericsConfig(
+    nexp=0,
+    minexp=0,
+    nmant=4,
+    has_subnormals=True,
+    has_two_nan=False,
+    has_naninf=False,
+)
+
+##########
+## FP4
+##########
 e3m0 = FpNumericsConfig(
     nexp=3,
     minexp=0,
@@ -121,6 +175,45 @@ class FpNumericsConfig:
     has_naninf=False,
 )
 
+e1m2 = FpNumericsConfig(
+    nexp=1,
+    minexp=0,
+    nmant=2,
+    has_subnormals=False,
+    has_two_nan=False,
+    has_naninf=False,
+)
+
+e1m2_ocp = FpNumericsConfig(
+    nexp=1,
+    minexp=0,
+    nmant=2,
+    has_subnormals=True,
+    has_two_nan=False,
+    has_naninf=False,
+)
+
+e0m3 = FpNumericsConfig(
+    nexp=0,
+    minexp=0,
+    nmant=3,
+    has_subnormals=False,
+    has_two_nan=False,
+    has_naninf=False,
+)
+
+e0m3_ocp = FpNumericsConfig(
+    nexp=0,
+    minexp=0,
+    nmant=3,
+    has_subnormals=True,
+    has_two_nan=False,
+    has_naninf=False,
+)
+
+##########
+## Misc
+##########
 e1m0 = FpNumericsConfig(
     nexp=1,
     minexp=0,