Repository: wandb/wandb

Issue:
[CLI]: `jnp.bfloat16` device arrays not supported by `wandb.log`
### Describe the bug

<!--- Description of the issue below  -->
Calling `wandb.log` on a dictionary containing jax.numpy (`jnp`) bfloat16 device arrays throws a recursion error.
<!--- A minimal code snippet between the quotes below  -->
```python 
import wandb
import jax.numpy as jnp

wandb.init(project="dummy_project")

metrics = {"loss": jnp.array(1., dtype=jnp.float32)}

wandb.log(metrics, step=1) # works fine!

metrics = {"loss": jnp.array(1., dtype=jnp.bfloat16)}

wandb.log(metrics, step=2) # recursion error
```

<!--- A full traceback of the exception in the quotes below -->
```shell
Traceback (most recent call last):
  File "test.py", line 12, in <module>
    wandb.log(metrics, step=2) # recursion error
  File "/Users/sanchitgandhi/venv/lib/python3.8/site-packages/wandb/sdk/wandb_run.py", line 1422, in log
    self._log(data=data, step=step, commit=commit)
  File "/Users/sanchitgandhi/venv/lib/python3.8/site-packages/wandb/sdk/wandb_run.py", line 1235, in _log
    self._partial_history_callback(data, self.history_step)
  File "/Users/sanchitgandhi/venv/lib/python3.8/site-packages/wandb/sdk/wandb_run.py", line 1083, in _partial_history_callback
    self._backend.interface.publish_partial_history(
  File "/Users/sanchitgandhi/venv/lib/python3.8/site-packages/wandb/sdk/interface/interface.py", line 511, in publish_partial_history
    item.value_json = json_dumps_safer_history(v)
  File "/Users/sanchitgandhi/venv/lib/python3.8/site-packages/wandb/util.py", line 810, in json_dumps_safer_history
    return json.dumps(obj, cls=WandBHistoryJSONEncoder, **kwargs)
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/json/__init__.py", line 234, in dumps
    return cls(
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/json/encoder.py", line 199, in encode
    chunks = self.iterencode(o, _one_shot=True)
  File "/Library/Developer/CommandLineTools/Library/Frameworks/Python3.framework/Versions/3.8/lib/python3.8/json/encoder.py", line 257, in iterencode
    return _iterencode(o, 0)
  File "/Users/sanchitgandhi/venv/lib/python3.8/site-packages/wandb/util.py", line 773, in default
    obj, converted = json_friendly(obj)
  File "/Users/sanchitgandhi/venv/lib/python3.8/site-packages/wandb/util.py", line 624, in json_friendly
    if getsizeof(obj) > VALUE_BYTES_LIMIT:
RecursionError: maximum recursion depth exceeded while calling a Python object
```


### Additional Files

_No response_

### Environment

WandB version:
0.12.11
OS:
MacOS (local) & Ubuntu 20.04 (gcloud)
Python version:
3.8
Versions of relevant libraries:
JAX 0.3.1

### Additional Context

_No response_


Discussion:
Leslie commented: 
Hi, I was able to run your code and repro this error. I'll put a ticket for it, but in the meantime you can use dtype=jnp.float32 since it works with our product.


Resolving patch:
diff --git a/wandb/util.py b/wandb/util.py
--- a/wandb/util.py
+++ b/wandb/util.py
@@ -607,11 +607,14 @@ def json_friendly(  # noqa: C901
         obj = obj.item()
         if isinstance(obj, float) and math.isnan(obj):
             obj = None
-        elif isinstance(obj, np.generic) and obj.dtype.kind == "f":
+        elif isinstance(obj, np.generic) and (
+            obj.dtype.kind == "f" or obj.dtype == "bfloat16"
+        ):
             # obj is a numpy float with precision greater than that of native python float
-            # (i.e., float96 or float128). in this case obj.item() does not return a native
-            # python float to avoid loss of precision, so we need to explicitly cast this
-            # down to a 64bit float
+            # (i.e., float96 or float128) or it is of custom type such as bfloat16.
+            # in these cases, obj.item() does not return a native
+            # python float (in the first case - to avoid loss of precision,
+            # so we need to explicitly cast this down to a 64bit float)
             obj = float(obj)
 
     elif isinstance(obj, bytes):


Test patch:
diff --git a/functional_tests/jax/01-log-bfloat16.py b/functional_tests/jax/01-log-bfloat16.py
new file mode 100644
--- /dev/null
+++ b/functional_tests/jax/01-log-bfloat16.py
@@ -0,0 +1,14 @@
+#!/usr/bin/env python
+"""Log a bfloat16 tensor."""
+
+import jax.numpy as jnp
+import wandb
+
+
+if __name__ == "__main__":
+    run = wandb.init()
+    m1 = jnp.array(1., dtype=jnp.float32)
+    m2 = jnp.array(2., dtype=jnp.bfloat16)
+    m3 = jnp.array([3., 4.], dtype=jnp.bfloat16)
+    run.log(dict(m1=m1, m2=m2, m3=m3))
+    run.finish()
diff --git a/functional_tests/jax/01-log-bfloat16.yea b/functional_tests/jax/01-log-bfloat16.yea
new file mode 100644
--- /dev/null
+++ b/functional_tests/jax/01-log-bfloat16.yea
@@ -0,0 +1,22 @@
+id: 0.jax.01-log-bfloat16
+tag:
+  platforms:
+    - linux
+    - mac
+plugin:
+  - wandb
+depend:
+  requirements:
+    - jax
+    - jaxlib
+    - numpy
+assert:
+  - :wandb:runs_len: 1
+  - :wandb:runs[0][config]: {}
+  - :wandb:runs[0][summary]:
+      m1: 1
+      m2: 2
+      m3:
+        - 3
+        - 4
+  - :wandb:runs[0][exitcode]: 0
diff --git a/functional_tests/xgboost/test_xgboost_clf.py b/functional_tests/xgboost/test_xgboost_clf.py
--- a/functional_tests/xgboost/test_xgboost_clf.py
+++ b/functional_tests/xgboost/test_xgboost_clf.py
@@ -11,16 +11,17 @@
 
 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1)
 
-model = XGBClassifier(use_label_encoder=False,
-                      eval_metric=['mlogloss', 'auc'],
-                      seed=42,
-                      n_estimators=50)
+model = XGBClassifier(
+    use_label_encoder=False, eval_metric=["mlogloss", "auc"], seed=42, n_estimators=50
+)
 
 wandb.init(project="wine-xgboost")
 
-model.fit(X_train,
-          y_train,
-          eval_set=[(X_train, y_train), (X_test, y_test)],
-          early_stopping_rounds=40,
-          callbacks=[WandbCallback(log_model=True)],
-          verbose=False)
+model.fit(
+    X_train,
+    y_train,
+    eval_set=[(X_train, y_train), (X_test, y_test)],
+    early_stopping_rounds=40,
+    callbacks=[WandbCallback(log_model=True)],
+    verbose=False,
+)
diff --git a/functional_tests/xgboost/test_xgboost_clf_old.py b/functional_tests/xgboost/test_xgboost_clf_old.py
--- a/functional_tests/xgboost/test_xgboost_clf_old.py
+++ b/functional_tests/xgboost/test_xgboost_clf_old.py
@@ -11,16 +11,17 @@
 
 X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1)
 
-model = XGBClassifier(use_label_encoder=False,
-                      eval_metric=['mlogloss', 'auc'],
-                      seed=42,
-                      n_estimators=50)
+model = XGBClassifier(
+    use_label_encoder=False, eval_metric=["mlogloss", "auc"], seed=42, n_estimators=50
+)
 
 wandb.init(project="wine-xgboost")
 
-model.fit(X_train,
-          y_train,
-          eval_set=[(X_train, y_train), (X_test, y_test)],
-          early_stopping_rounds=40,
-          callbacks=[wandb_callback()],
-          verbose=False)
+model.fit(
+    X_train,
+    y_train,
+    eval_set=[(X_train, y_train), (X_test, y_test)],
+    early_stopping_rounds=40,
+    callbacks=[wandb_callback()],
+    verbose=False,
+)
diff --git a/functional_tests/xgboost/test_xgboost_reg.py b/functional_tests/xgboost/test_xgboost_reg.py
--- a/functional_tests/xgboost/test_xgboost_reg.py
+++ b/functional_tests/xgboost/test_xgboost_reg.py
@@ -13,21 +13,23 @@
 # load data
 housing = fetch_california_housing()
 data = pd.DataFrame(housing.data)
-X, y = data.iloc[:,:-1],data.iloc[:,-1]
-data_dmatrix = xgb.DMatrix(data=X,label=y)
+X, y = data.iloc[:, :-1], data.iloc[:, -1]
+data_dmatrix = xgb.DMatrix(data=X, label=y)
 
 # Train validation split
-X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=123)
+X_train, X_test, y_train, y_test = train_test_split(
+    X, y, test_size=0.2, random_state=123
+)
 
 # Define regressor
 bst_params = dict(
-    objective='reg:squarederror',
+    objective="reg:squarederror",
     colsample_bytree=0.3,
     learning_rate=0.1,
     max_depth=5,
     alpha=10,
     n_estimators=100,
-    tree_method='hist'
+    tree_method="hist",
 )
 
 xg_reg = xgb.XGBRegressor(**bst_params)
@@ -35,11 +37,14 @@
 # Initialize run
 wandb.init(project="xgboost-housing")
 
-xg_reg.fit(X_train, y_train,
-           eval_set=[(X_train, y_train), (X_test, y_test)],
-           early_stopping_rounds=20,
-           callbacks=[WandbCallback()],
-           verbose=False)
+xg_reg.fit(
+    X_train,
+    y_train,
+    eval_set=[(X_train, y_train), (X_test, y_test)],
+    early_stopping_rounds=20,
+    callbacks=[WandbCallback()],
+    verbose=False,
+)
 
 # Evaluate
 preds = xg_reg.predict(X_test)
diff --git a/functional_tests/xgboost/xgboost_clf_old.yea b/functional_tests/xgboost/xgboost_clf_old.yea
--- a/functional_tests/xgboost/xgboost_clf_old.yea
+++ b/functional_tests/xgboost/xgboost_clf_old.yea
@@ -5,7 +5,7 @@ command:
     program: test_xgboost_clf_old.py
 depend:
     requirements:
-        - xgboost>=1.3.0
+        - xgboost>=1.3.0,<1.6.0
         
 assert:
   - :wandb:runs_len: 1
diff --git a/tests/test_util.py b/tests/test_util.py
--- a/tests/test_util.py
+++ b/tests/test_util.py
@@ -506,3 +506,21 @@ def test_resolve_aliases():
 
     aliases = util._resolve_aliases("boom")
     assert aliases == ["boom", "latest"]
+
+
+@pytest.mark.skipif(
+    platform.system() == "Windows", reason="test suite does not build jaxlib on windows"
+)
+def test_bfloat16_to_float():
+    import jax.numpy as jnp
+
+    array = jnp.array(1.0, dtype=jnp.bfloat16)
+    # array to scalar bfloat16
+    array_cast = util.json_friendly(array)
+    assert array_cast[1] is True
+    assert array_cast[0].__class__.__name__ == "bfloat16"
+    # scalar bfloat16 to float
+    array_cast = util.json_friendly(array_cast[0])
+    assert array_cast[0] == 1.0
+    assert array_cast[1] is True
+    assert isinstance(array_cast[0], float)