(bf16_arr: np.ndarray[Any, np.dtype[np.uint16]])
| 716 | |
| 717 | |
| 718 | def bf16_to_fp32(bf16_arr: np.ndarray[Any, np.dtype[np.uint16]]) -> NDArray: |
| 719 | assert bf16_arr.dtype == np.uint16, f"Input array should be of dtype uint16, but got {bf16_arr.dtype}" |
| 720 | fp32_arr = bf16_arr.astype(np.uint32) << 16 |
| 721 | return fp32_arr.view(np.float32) |
| 722 | |
| 723 | def preprocess_weights( |
| 724 | w: np.ndarray, |