bfloat16 @ bfloat16 -> float32?

I see different behaviors in NumPy and JAX.

In NumPy,

```py
>>> import numpy as np
>>> a=np.ones((4,4), dtype=ml_dtypes.bfloat16)
>>> a@a
array([[4., 4., 4., 4.],
[4., 4., 4., 4.],
[4., 4., 4., 4.],
[4., 4., 4., 4.]], dtype=float32)
```
The data type of the output is `float32`.

In JAX,

```py
>>> import jax.numpy as jnp
>>> b=jnp.asarray(a)
>>> b@b
Array([[4, 4, 4, 4],
[4, 4, 4, 4],
[4, 4, 4, 4],
[4, 4, 4, 4]], dtype=bfloat16)
```

The data type of the output is `bfloat16`.

I need clarification about this behavior. I don't see any documentation about it, so I am unsure whether it is a bug or a feature.