Repository: google/flax

Issue:
Incompatible variables for Tensorboard hparams are recast to strings but never returned
### Core Problem
Tensorboard hparams only supports a subset of Python and Numpy variable types ([see hparams docstrings](https://github.com/tensorflow/tensorboard/blob/1204566da5437af55109f7a4af18f9f8b7c4f864/tensorboard/plugins/hparams/summary_v2.py)). The `flax.metrics.tensorboard.SummaryWriter` class's method `SummaryWriter.hparams()` should handle this behavior via the `flax.metrics.tensorboard._flatten_dict()` function, casting incompatible types to strings (which hparams supports). However, despite performing the casting operation, the `_flatten_dict` function does not append the recast variables to the dictionary it returns. 

The result, for the below example, is that the "hidden_layers" parameters are silently excluded and do not appear in Tensorboard's hparams.

```Python 
from flax.metrics import tensorboard

experiment_dir = "./Example"

network_hyperparameters = {
    "hidden_layers_list": [12,12],
    "hidden_layers_tuple": (12,12),
    "dropout_rate": 1.0,
}

summary_writer = tensorboard.SummaryWriter(experiment_dir)
summary_writer.hparams(network_hyperparameters)
summary_writer.scalar('Training loss', 0.1, 1)
summary_writer.flush()
```

### Colab Example:
[Example notebook](https://colab.research.google.com/gist/tttc3/8dd7ef04c4222bc18fb03b043d370120/falx_tensorboard_issue_demo.ipynb)

### Proposed fix
Modify `_flattened_dict` to explicitly check if a dictionary value is one of those supported by Tensorboard's hparams api, as defined [here](https://github.com/tensorflow/tensorboard/blob/1204566da5437af55109f7a4af18f9f8b7c4f864/tensorboard/plugins/hparams/summary_v2.py). If the value is not supported, cast it to a string and append it to the dictionary that `_flattened_dict` normally returns.

**Current _flatten_dict code**
```Python
def _flatten_dict(input_dict, parent_key='', sep='.'):
  """Flattens and simplifies dict such that it can be used by hparams.

  Args:
    input_dict: Input dict, e.g., from ConfigDict.
    parent_key: String used in recursion.
    sep: String used to separate parent and child keys.

  Returns:
   Flattened dict.
  """
  items = []
  for k, v in input_dict.items():
    new_key = parent_key + sep + k if parent_key else k

    # Take special care of things hparams cannot handle.
    if v is None:
      v = 'None'
    elif isinstance(v, list):
      v = str(v)
    elif isinstance(v, tuple):
      v = str(v)
    elif isinstance(v, dict):
      # Recursively flatten the dict.
      items.extend(_flatten_dict(v, new_key, sep=sep).items())
    else:
      items.append((new_key, v))
  return dict(items)
```

**Proposed _flatten_dict code modification**
```Python
def _flatten_dict(input_dict, parent_key='', sep='.'):
  """Flattens and simplifies dict such that it can be used by hparams.

  Args:
    input_dict: Input dict, e.g., from ConfigDict.
    parent_key: String used in recursion.
    sep: String used to separate parent and child keys.

  Returns:
   Flattened dict.
  """
  items = []
  for k, v in input_dict.items():
    new_key = parent_key + sep + k if parent_key else k

    # Valid types according to https://github.com/tensorflow/tensorboard/blob/1204566da5437af55109f7a4af18f9f8b7c4f864/tensorboard/plugins/hparams/summary_v2.py
    valid_types = (bool, int, float, str, np.bool_, np.integer, np.floating, np.character)

    if isinstance(v, dict):
      # Recursively flatten the dict.
      items.extend(_flatten_dict(v, new_key, sep=sep).items())
      continue
    elif not isinstance(v, valid_types):
      # Cast any incompatible values as strings such that they can be handled by hparams
      v = str(v)
    items.append((new_key, v))
  return dict(items)
```

I am happy submit a pull request with the modifications.


Discussion:
Thanks for noticing this. Indeed there seems to be a bug in our code, and we actually do nothing with `v` if it is `None`, `list` or `tuple`! Yes, it would be great if you could file this as a PR and I think your suggested change using `valid_types` is an improvement.

We should also run internals tests on this to make sure your change doesn't break anything.

Resolving patch:
diff --git a/flax/metrics/tensorboard.py b/flax/metrics/tensorboard.py
--- a/flax/metrics/tensorboard.py
+++ b/flax/metrics/tensorboard.py
@@ -39,18 +39,17 @@ def _flatten_dict(input_dict, parent_key='', sep='.'):
   for k, v in input_dict.items():
     new_key = parent_key + sep + k if parent_key else k
 
-    # Take special care of things hparams cannot handle.
-    if v is None:
-      v = 'None'
-    elif isinstance(v, list):
-      v = str(v)
-    elif isinstance(v, tuple):
-      v = str(v)
-    elif isinstance(v, dict):
+    # Valid types according to https://github.com/tensorflow/tensorboard/blob/1204566da5437af55109f7a4af18f9f8b7c4f864/tensorboard/plugins/hparams/summary_v2.py
+    valid_types = (bool, int, float, str, np.bool_, np.integer, np.floating, np.character)
+
+    if isinstance(v, dict):
       # Recursively flatten the dict.
       items.extend(_flatten_dict(v, new_key, sep=sep).items())
-    else:
-      items.append((new_key, v))
+      continue
+    elif not isinstance(v, valid_types):
+      # Cast any incompatible values as strings such that they can be handled by hparams
+      v = str(v)
+    items.append((new_key, v))
   return dict(items)
 
 


Test patch:
diff --git a/tests/tensorboard_test.py b/tests/tensorboard_test.py
--- a/tests/tensorboard_test.py
+++ b/tests/tensorboard_test.py
@@ -24,7 +24,7 @@
 from tensorboard.util import tensor_util
 import tensorflow.compat.v2 as tf
 
-from flax.metrics.tensorboard import SummaryWriter
+from flax.metrics.tensorboard import SummaryWriter, _flatten_dict
 
 def _process_event(event):
   for value in event.summary.value:
@@ -262,5 +262,58 @@ def test_summarywriter_histogram_2bins(self):
     self.assertTrue(
         np.allclose(actual_histogram[1], (499.5, 999.0, 500.0), atol=1e-01))
 
+  def test_flatten_dict(self):
+    # Valid types according to https://github.com/tensorflow/tensorboard/blob/1204566da5437af55109f7a4af18f9f8b7c4f864/tensorboard/plugins/hparams/summary_v2.py
+    input_hparams={
+      # Example Invalid Types
+      "None": None, "List": [1, 2, 3], "Tuple": (1, 2, 3), "Complex": complex("1+1j"), "np.complex_": np.complex_("1+1j"),
+      # Valid Python Types
+      "Bool": True, "Int": 1, "Float": 1.0, "Str": "test",
+      # Valid Numpy Types
+      "np.bool_": np.bool_(1), "np.integer": np.int_(1),  "np.floating": np.float_(1.0), "np.character": np.str_("test"),
+      # Nested dict to flatten
+      "Nested_Dict": {
+        "None": None,
+        "List": [1, 2, 3],
+        "Tuple": (1, 2, 3),
+        "Complex": complex("1+1j"), 
+        "np.complex_": np.complex_("1+1j"),
+        "Bool": True,
+        "Int": 1,
+        "Float": 1.0,
+        "Str": "test",
+        "np.bool_": np.bool_(1),
+        "np.integer": np.int_(1),
+        "np.floating": np.float_(1.0),
+        "np.character": np.str_("test")
+      }
+    }
+    
+    result_hparams = _flatten_dict(input_hparams)
+
+    expected_hparams={
+      "None": "None", "List": "[1, 2, 3]", "Tuple": "(1, 2, 3)", "Complex": "(1+1j)", "np.complex_": "(1+1j)",
+      # Valid Python Types
+      "Bool": True, "Int": 1, "Float": 1.0, "Str": "test",
+      # Valid Numpy Types
+      "np.bool_": np.bool_(1), "np.integer": np.int_(1),  "np.floating": np.float_(1.0), "np.character": np.str_("test"),
+      # Nested Dict
+      "Nested_Dict.None": "None",
+      "Nested_Dict.List": "[1, 2, 3]",
+      "Nested_Dict.Tuple": "(1, 2, 3)",
+      "Nested_Dict.Complex": "(1+1j)",
+      "Nested_Dict.np.complex_": "(1+1j)",
+      "Nested_Dict.Bool": True,
+      "Nested_Dict.Int": 1,
+      "Nested_Dict.Float": 1.0,
+      "Nested_Dict.Str": "test",
+      "Nested_Dict.np.bool_": np.bool_(1),
+      "Nested_Dict.np.integer": np.int_(1),
+      "Nested_Dict.np.floating": np.float_(1.0),
+      "Nested_Dict.np.character": np.str_("test")
+    }
+
+    self.assertDictEqual(result_hparams, expected_hparams)
+
 if __name__ == '__main__':
   absltest.main()