Allow bfloat16 default dtype (#19074)

Useful for llms! The tradeoff in precision can often be worth it in
memory constrained environments, and unlike float16, does not have
the same overflow/underflow issues during training.

keras/backend/config.py

@@ -20,7 +20,7 @@
 def floatx():
  """Return the default float type, as a string.
 
- E.g. `'float16'`, `'float32'`, `'float64'`.
+ E.g. `'bfloat16'`, `'float16'`, `'float32'`, `'float64'`.
 
  Returns:
  String, the current default float type.
@@ -45,7 +45,7 @@ def set_floatx(value):
  `keras.mixed_precision.set_dtype_policy('mixed_float16')`.
 
  Args:
- value: String; `'float16'`, `'float32'`, or `'float64'`.
+ value: String; `'bfloat16'`, `'float16'`, `'float32'`, or `'float64'`.
 
  Examples:
  >>> keras.config.floatx()
@@ -62,7 +62,7 @@ def set_floatx(value):
  ValueError: In case of invalid value.
  """
  global _FLOATX
- accepted_dtypes = {"float16", "float32", "float64"}
+ accepted_dtypes = {"bfloat16", "float16", "float32", "float64"}
  if value not in accepted_dtypes:
  raise ValueError(
  f"Unknown `floatx` value: {value}. "