jax-ml / jax-ml/ml_dtypes

bfloat16 & other custom dtypes incompatible with np.save

Open
#41 1 comment 0 reactions 0 assignees View on GitHub
bug
Dominant language
C++
Stars
361
Forks
60
Avg merge
18h 41m
Merged PRs (30d)
3

Description

```python
from ml_dtypes import bfloat16
import numpy as np

x = np.arange(4, dtype=bfloat16)
np.save('out.npy', x)
print(np.load('out.npy'))
# [b'\x00\x00' b'\x80\x3F' b'\x00\x40' b'\x40\x40']
```
This is due to the fact that `np.save` saves the dtype via the `kind` code, and we chose `np.dtype(bfloat16).kind = 'V'` because such codes are not extensible in numpy.

When reconstructing the dtype, this is the result:
```python
np.dtype(np.dtype(bfloat16).kind)
# dtype('V')
```
I don't think this is something that can be fixed without deeper changes to numpy itself.

Contributor guide

Open the contributing guide

Research direction

Start by reproducing the example with ml_dtypes.bfloat16, np.save, and np.load, then inspect how the dtype kind is serialized and reconstructed. There are no files or tests named in the issue; done would require an agreed solution for preserving custom dtypes through NumPy’s save/load path, potentially involving deeper NumPy changes.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.