ENH: Native bfloat16 (bf16) dtype support for unified memory AI workloads

[bkpaine1]

Proposed new feature or change:

Summary

As unified memory architectures become mainstream for AI inference (AMD APUs, Apple Silicon, NVIDIA Grace Hopper, upcoming DGX Spark), the lack of native bfloat16 support in numpy is causing systematic failures at GPU↔CPU boundaries.

The Problem

Modern AI models output bfloat16 tensors for memory efficiency. On discrete GPUs, the explicit .cpu() copy across PCIe often triggers implicit dtype conversion, masking the issue.

On unified memory architectures, GPU and CPU share the same memory space - no copy occurs. When code calls .numpy() on a bf16 tensor:

# Works on discrete GPU (implicit conversion during PCIe copy)
# Crashes on unified memory (tensor already visible to CPU)
waveform = audio_tensor.squeeze(0).numpy()

Error:

TypeError: Got unsupported ScalarType BFloat16

Affected Architectures

This impacts ALL unified memory systems:

Platform	Status	Use Case
AMD Strix Halo APU	Shipping now	128GB unified, AI workstations
Apple Silicon M1-M4	Shipping	Mac AI/ML workflows
NVIDIA Grace Hopper	Shipping	Datacenter AI
NVIDIA DGX Spark	Coming 2025	Desktop AI workstation
Intel Arc w/ unified mem	Shipping	Consumer AI

Current Workaround

Every GPU→numpy boundary requires explicit conversion:

# Before (breaks on unified memory)
array = tensor.numpy()

# After (works everywhere)
array = tensor.float().cpu().numpy()

This requires patching dozens of libraries: torchaudio, librosa, soundfile integrations, video I/O, etc.

Real-World Impact

In a single ComfyUI AI video pipeline, I've had to patch:

• torchaudio resampler
• VibeVoice TTS output
• LatentSync lip-sync
• Wav2Lip lip-sync
• Video frame I/O

Each library assumes numpy can handle whatever dtype PyTorch gives it. On unified memory, this assumption fails.

Proposed Solution

Add native bfloat16 dtype to numpy:

import numpy as np

# Proposed
arr = np.array([1.0, 2.0, 3.0], dtype=np.bfloat16)

# Or at minimum, automatic conversion on array creation
bf16_tensor = torch.tensor([1.0], dtype=torch.bfloat16, device='cuda')
arr = bf16_tensor.numpy()  # Should auto-convert to float32 with warning

Alternatives Considered

1. Do nothing - Force every AI library to add .float() calls (current state, not scalable)
2. Warn and convert - numpy accepts bf16, warns, converts to float32 automatically
3. Full support - Native bf16 dtype with arithmetic operations

Option 2 would solve 90% of the pain with minimal implementation effort.

Context

bfloat16 was designed by Google Brain specifically for ML workloads. It's now the default output dtype for:

• Hugging Face Transformers
• PyTorch model inference
• Most modern AI models

As unified memory becomes the standard for AI hardware (it's more efficient), this gap between PyTorch's bf16 support and numpy's lack thereof will cause increasing friction.

Environment

• numpy: 1.26.x / 2.x
• PyTorch: 2.5+ with ROCm/CUDA
• Hardware: AMD Strix Halo (gfx1151), 128GB unified memory
• OS: Ubuntu 24.04

References

• PyTorch bfloat16: https://pytorch.org/docs/stable/tensors.html
• Google Brain bfloat16: https://en.wikipedia.org/wiki/Bfloat16_floating-point_format
• AMD ROCm unified memory: https://rocm.docs.amd.com/
• NVIDIA Grace Hopper unified memory: https://www.nvidia.com/en-us/data-center/grace-hopper-superchip/

---

The unified memory future is here. numpy's dtype support should reflect that reality.

[mattip]

NumPy supports externally defined dtypes, and bfloat16 is specifically discussed in #19808. That issue led to numpy-user-dtypes and also points to bfloat16. Converting to np.float32 or np.half is probably not what you want to do.

[rkern]

Alternatives Considered

1. Do nothing - Force every AI library to add .float() calls (current state, not scalable)
2. Warn and convert - numpy accepts bf16, warns, converts to float32 automatically
3. Full support - Native bf16 dtype with arithmetic operations

Option 2 would solve 90% of the pain with minimal implementation effort.

Seems to me like Option 2 would be the way to go if that tradeoff were the case. However, the implementation isn't ours to do. The TypeError: Got unsupported ScalarType BFloat16 is from pytorch. It's their .numpy() method. If we implement a bfloat16 dtype in numpy, pytorch's updated .numpy() method would create bfloat16 numpy arrays and return those. We would never get the chance to "warn and convert to float32 automatically". pytorch, on the other hand, could special-case BFloat16 earlier and do the internal equivalent of .float().cpu(), with a warning, before passing it along to the conversion to ndarrays.

pytorch could also adopt the 3rd party bfloat16 implementation linked above either as a dependency or vendor it in. I'd like to hear from them. I think this request is more about PyTorch than us, at this point. Now, if they say that they have reasons such that they only support it if it were a builtin numpy dtype, that'd be another thing. Or better, if you got 2 Array API implementors coming along to say that they both need it as a builtin numpy dtype.

[ngoldbaum]

bfloat16 is also provided by the ml_dtypes project maintained by the jax team: https://github.com/jax-ml/ml_dtypes

[bkpaine1]

bfloat16 is also provided by the ml_dtypes project maintained by the jax team: https://github.com/jax-ml/ml_dtypes

I am just fighting my way through it bit by bit with this Halo Strix. TBH I am so deep in the weeds but have rarely open issues or comments, forgive my learning curve how to do this. But I do fix it in numpy. that is where I put the patch to make things work. bf float works for video render but not audio weird. Still just pushing numpy to cpu fixed it. I guess my question is are you saying I should open an issue with the ml guys too?

[rkern]

But I do fix it in numpy. that is where I put the patch to make things work.

What patch? Which "numpy" do you mean? The numpy package? Or the .numpy() method in torch? If the latter, yeah, you should probably be talking to PyTorch.

[bkpaine1]

But I do fix it in numpy. that is where I put the patch to make things work.

What patch? Which "numpy" do you mean? The numpy package? Or the .numpy() method in torch? If the latter, yeah, you should probably be talking to PyTorch.

I'm patching the application code that calls .numpy() - adding .float() before the numpy call. Not patching numpy or torch packages themselves. The proper fix would be numpy adding native bfloat16 dtype support so the cast isn't needed."
The fix SHOULD go in numpy (add bfloat16 support), but I am WORKING AROUND it in downstream applications. the ask for numpy to accept the bfloat16. and not torch several packages in comfuyi and others. the solve would in numpy.

[rgommers]

I'll close this as a duplicate of gh-19808.