Look for an existing bfloat16 type and prefer its use.
Look for an existing bfloat16 type and prefer its use.

This is a stop gap measure to allow JAX and TensorFlow to share the same bfloat16 NumPy type in opensource, where they are built as separate .so files.

A better fix would be to split the bfloat16 extension out of the TensorFlow tree and into its own pip package, but that is more work. For now, change both JAX and TF to prefer an existing NumPy "bfloat16" type over their own version, if one is registered.

PiperOrigin-RevId: 346619498

google/jax,google/jax,google/jax,google/jax