Returns (mantissa_bits, total_bits) for each format. mantissa_bits excludes the implicit leading 1.
(dtype: torch.dtype)
| 119 | |
| 120 | |
| 121 | def get_format_params(dtype: torch.dtype) -> tuple[int, int]: |
| 122 | """ |
| 123 | Returns (mantissa_bits, total_bits) for each format. |
| 124 | mantissa_bits excludes the implicit leading 1. |
| 125 | """ |
| 126 | if dtype == torch.float32: |
| 127 | return 23, 32 |
| 128 | elif dtype == torch.bfloat16: |
| 129 | return 7, 16 |
| 130 | elif dtype == torch.float16: |
| 131 | return 10, 16 |
| 132 | elif dtype == torch.float8_e4m3fn: |
| 133 | return 3, 8 |
| 134 | elif dtype == torch.float8_e5m2: |
| 135 | return 2, 8 |
| 136 | elif dtype == torch.int8: |
| 137 | return 0, 8 # Int8 doesn't have mantissa bits |
| 138 | else: |
| 139 | raise ValueError(f"Unsupported dtype: {dtype}") |
| 140 | |
| 141 | |
| 142 | def copy_stochastic( |