We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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.

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `flax/metrics/tensorboard.py`
Content:
```
1 # Copyright 2022 The Flax Authors.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 """Write Summaries from JAX for use with Tensorboard.
16 """
17 
18 import os
19 
20 # pylint: disable=g-import-not-at-top
21 import numpy as np
22 
23 import tensorflow.compat.v2 as tf
24 from tensorboard.plugins.hparams import api as hparams_api
25 
26 
27 def _flatten_dict(input_dict, parent_key='', sep='.'):
28   """Flattens and simplifies dict such that it can be used by hparams.
29 
30   Args:
31     input_dict: Input dict, e.g., from ConfigDict.
32     parent_key: String used in recursion.
33     sep: String used to separate parent and child keys.
34 
35   Returns:
36    Flattened dict.
37   """
38   items = []
39   for k, v in input_dict.items():
40     new_key = parent_key + sep + k if parent_key else k
41 
42     # Take special care of things hparams cannot handle.
43     if v is None:
44       v = 'None'
45     elif isinstance(v, list):
46       v = str(v)
47     elif isinstance(v, tuple):
48       v = str(v)
49     elif isinstance(v, dict):
50       # Recursively flatten the dict.
51       items.extend(_flatten_dict(v, new_key, sep=sep).items())
52     else:
53       items.append((new_key, v))
54   return dict(items)
55 
56 
57 class SummaryWriter(object):
58   """Saves data in event and summary protos for tensorboard."""
59 
60   def __init__(self, log_dir):
61     """Create a new SummaryWriter.
62 
63     Args:
64       log_dir: path to record tfevents files in.
65     """
66     log_dir = os.fspath(log_dir)
67 
68     # If needed, create log_dir directory as well as missing parent directories.
69     if not tf.io.gfile.isdir(log_dir):
70       tf.io.gfile.makedirs(log_dir)
71 
72     self._event_writer = tf.summary.create_file_writer(log_dir, 10, 120, None)
73     self._closed = False
74 
75   def close(self):
76     """Close SummaryWriter. Final!"""
77     if not self._closed:
78       self._event_writer.close()
79       self._closed = True
80       del self._event_writer
81 
82   def flush(self):
83     self._event_writer.flush()
84 
85   def scalar(self, tag, value, step):
86     """Saves scalar value.
87 
88     Args:
89       tag: str: label for this data
90       value: int/float: number to log
91       step: int: training step
92     """
93     value = float(np.array(value))
94     with self._event_writer.as_default():
95       tf.summary.scalar(name=tag, data=value, step=step)
96 
97   def image(self, tag, image, step, max_outputs=3):
98     """Saves RGB image summary from np.ndarray [H,W], [H,W,1], or [H,W,3].
99 
100     Args:
101       tag: str: label for this data
102       image: ndarray: [H,W], [H,W,1], [H,W,3], [K,H,W], [K,H,W,1], [K,H,W,3]
103         Save image in greyscale or colors.
104         Pixel values could be either uint8 or float.
105         Floating point values should be in range [0, 1).
106       step: int: training step
107       max_outputs: At most this many images will be emitted at each step.
108         Defaults to 3.
109     """
110     image = np.array(image)
111     # tf.summary.image expects image to have shape [k, h, w, c] where,
112     # k = number of samples, h = height, w = width, c = number of channels.
113     if len(np.shape(image)) == 2:
114       image = image[np.newaxis, :, :, np.newaxis]
115     elif len(np.shape(image)) == 3:
116       # this could be either [k, h, w] or [h, w, c]
117       if np.shape(image)[-1] in (1, 3):
118         image = image[np.newaxis, :, :, :]
119       else:
120         image = image[:, :, :, np.newaxis]
121     if np.shape(image)[-1] == 1:
122       image = np.repeat(image, 3, axis=-1)
123 
124     # Convert to tensor value as tf.summary.image expects data to be a tensor.
125     image = tf.convert_to_tensor(image)
126     with self._event_writer.as_default():
127       tf.summary.image(name=tag, data=image, step=step, max_outputs=max_outputs)
128 
129   def audio(self, tag, audiodata, step, sample_rate=44100, max_outputs=3):
130     """Saves audio as wave.
131 
132     NB: single channel only right now.
133 
134     Args:
135       tag: str: label for this data
136       audiodata: ndarray [Nsamples, Nframes, Nchannels]: audio data to
137         be saved as wave. The data will be clipped to [-1.0, 1.0].
138       step: int: training step
139       sample_rate: sample rate of passed in audio buffer
140       max_outputs: At most this many audio clips will be emitted at each
141         step. Defaults to 3.
142     """
143     # tf.summary.audio expects the audio data to have floating values in
144     # [-1.0, 1.0].
145     audiodata = np.clip(np.array(audiodata), -1, 1)
146 
147     # Convert to tensor value as tf.summary.audio expects data to be a tensor.
148     audio = tf.convert_to_tensor(audiodata, dtype=tf.float32)
149     with self._event_writer.as_default():
150       tf.summary.audio(
151           name=tag, data=audio, sample_rate=sample_rate, step=step,
152           max_outputs=max_outputs, encoding='wav')
153 
154   def histogram(self, tag, values, step, bins=None):
155     """Saves histogram of values.
156 
157     Args:
158       tag: str: label for this data
159       values: ndarray: will be flattened by this routine
160       step: int: training step
161       bins: number of bins in histogram
162     """
163     values = np.array(values)
164     values = np.reshape(values, -1)
165     with self._event_writer.as_default():
166       tf.summary.histogram(name=tag, data=values, step=step, buckets=bins)
167 
168   def text(self, tag, textdata, step):
169     """Saves a text summary.
170 
171     Args:
172       tag: str: label for this data
173       textdata: string
174       step: int: training step
175     Note: markdown formatting is rendered by tensorboard.
176     """
177     if not isinstance(textdata, (str, bytes)):
178       raise ValueError('`textdata` should be of the type `str` or `bytes`.')
179     with self._event_writer.as_default():
180       tf.summary.text(name=tag, data=tf.constant(textdata), step=step)
181 
182   def write(self, tag, tensor, step, metadata=None):
183     """Saves a arbitrary tensor summary.
184 
185     Useful when working with custom plugins or constructing a summary directly.
186 
187     Args:
188       tag: str: label for this data
189       tensor: ndarray: tensor data to save.
190       step: int: training step
191       metadata: Optional SummaryMetadata, as a proto or serialized bytes.
192     Note: markdown formatting is rendered by tensorboard.
193     """
194     with self._event_writer.as_default():
195       tf.summary.write(
196           tag=tag,
197           tensor=tensor,
198           step=step,
199           metadata=metadata)
200 
201   def hparams(self, hparams):
202     """Saves hyper parameters.
203 
204     Args:
205       hparams: Flat mapping from hyper parameter name to value.
206     """
207 
208     with self._event_writer.as_default():
209       hparams_api.hparams(hparams=_flatten_dict(hparams))
210
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
   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)

{"golden_diff": "diff --git a/flax/metrics/tensorboard.py b/flax/metrics/tensorboard.py\n--- a/flax/metrics/tensorboard.py\n+++ b/flax/metrics/tensorboard.py\n@@ -39,18 +39,17 @@\n   for k, v in input_dict.items():\n     new_key = parent_key + sep + k if parent_key else k\n \n-    # Take special care of things hparams cannot handle.\n-    if v is None:\n-      v = 'None'\n-    elif isinstance(v, list):\n-      v = str(v)\n-    elif isinstance(v, tuple):\n-      v = str(v)\n-    elif isinstance(v, dict):\n+    # Valid types according to https://github.com/tensorflow/tensorboard/blob/1204566da5437af55109f7a4af18f9f8b7c4f864/tensorboard/plugins/hparams/summary_v2.py\n+    valid_types = (bool, int, float, str, np.bool_, np.integer, np.floating, np.character)\n+\n+    if isinstance(v, dict):\n       # Recursively flatten the dict.\n       items.extend(_flatten_dict(v, new_key, sep=sep).items())\n-    else:\n-      items.append((new_key, v))\n+      continue\n+    elif not isinstance(v, valid_types):\n+      # Cast any incompatible values as strings such that they can be handled by hparams\n+      v = str(v)\n+    items.append((new_key, v))\n   return dict(items)\n", "issue": "Incompatible variables for Tensorboard hparams are recast to strings but never returned\n### Core Problem\r\nTensorboard 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. \r\n\r\nThe result, for the below example, is that the \"hidden_layers\" parameters are silently excluded and do not appear in Tensorboard's hparams.\r\n\r\n```Python \r\nfrom flax.metrics import tensorboard\r\n\r\nexperiment_dir = \"./Example\"\r\n\r\nnetwork_hyperparameters = {\r\n    \"hidden_layers_list\": [12,12],\r\n    \"hidden_layers_tuple\": (12,12),\r\n    \"dropout_rate\": 1.0,\r\n}\r\n\r\nsummary_writer = tensorboard.SummaryWriter(experiment_dir)\r\nsummary_writer.hparams(network_hyperparameters)\r\nsummary_writer.scalar('Training loss', 0.1, 1)\r\nsummary_writer.flush()\r\n```\r\n\r\n### Colab Example:\r\n[Example notebook](https://colab.research.google.com/gist/tttc3/8dd7ef04c4222bc18fb03b043d370120/falx_tensorboard_issue_demo.ipynb)\r\n\r\n### Proposed fix\r\nModify `_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.\r\n\r\n**Current _flatten_dict code**\r\n```Python\r\ndef _flatten_dict(input_dict, parent_key='', sep='.'):\r\n  \"\"\"Flattens and simplifies dict such that it can be used by hparams.\r\n\r\n  Args:\r\n    input_dict: Input dict, e.g., from ConfigDict.\r\n    parent_key: String used in recursion.\r\n    sep: String used to separate parent and child keys.\r\n\r\n  Returns:\r\n   Flattened dict.\r\n  \"\"\"\r\n  items = []\r\n  for k, v in input_dict.items():\r\n    new_key = parent_key + sep + k if parent_key else k\r\n\r\n    # Take special care of things hparams cannot handle.\r\n    if v is None:\r\n      v = 'None'\r\n    elif isinstance(v, list):\r\n      v = str(v)\r\n    elif isinstance(v, tuple):\r\n      v = str(v)\r\n    elif isinstance(v, dict):\r\n      # Recursively flatten the dict.\r\n      items.extend(_flatten_dict(v, new_key, sep=sep).items())\r\n    else:\r\n      items.append((new_key, v))\r\n  return dict(items)\r\n```\r\n\r\n**Proposed _flatten_dict code modification**\r\n```Python\r\ndef _flatten_dict(input_dict, parent_key='', sep='.'):\r\n  \"\"\"Flattens and simplifies dict such that it can be used by hparams.\r\n\r\n  Args:\r\n    input_dict: Input dict, e.g., from ConfigDict.\r\n    parent_key: String used in recursion.\r\n    sep: String used to separate parent and child keys.\r\n\r\n  Returns:\r\n   Flattened dict.\r\n  \"\"\"\r\n  items = []\r\n  for k, v in input_dict.items():\r\n    new_key = parent_key + sep + k if parent_key else k\r\n\r\n    # Valid types according to https://github.com/tensorflow/tensorboard/blob/1204566da5437af55109f7a4af18f9f8b7c4f864/tensorboard/plugins/hparams/summary_v2.py\r\n    valid_types = (bool, int, float, str, np.bool_, np.integer, np.floating, np.character)\r\n\r\n    if isinstance(v, dict):\r\n      # Recursively flatten the dict.\r\n      items.extend(_flatten_dict(v, new_key, sep=sep).items())\r\n      continue\r\n    elif not isinstance(v, valid_types):\r\n      # Cast any incompatible values as strings such that they can be handled by hparams\r\n      v = str(v)\r\n    items.append((new_key, v))\r\n  return dict(items)\r\n```\r\n\r\nI am happy submit a pull request with the modifications.\n", "before_files": [{"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Write Summaries from JAX for use with Tensorboard.\n\"\"\"\n\nimport os\n\n# pylint: disable=g-import-not-at-top\nimport numpy as np\n\nimport tensorflow.compat.v2 as tf\nfrom tensorboard.plugins.hparams import api as hparams_api\n\n\ndef _flatten_dict(input_dict, parent_key='', sep='.'):\n  \"\"\"Flattens and simplifies dict such that it can be used by hparams.\n\n  Args:\n    input_dict: Input dict, e.g., from ConfigDict.\n    parent_key: String used in recursion.\n    sep: String used to separate parent and child keys.\n\n  Returns:\n   Flattened dict.\n  \"\"\"\n  items = []\n  for k, v in input_dict.items():\n    new_key = parent_key + sep + k if parent_key else k\n\n    # Take special care of things hparams cannot handle.\n    if v is None:\n      v = 'None'\n    elif isinstance(v, list):\n      v = str(v)\n    elif isinstance(v, tuple):\n      v = str(v)\n    elif isinstance(v, dict):\n      # Recursively flatten the dict.\n      items.extend(_flatten_dict(v, new_key, sep=sep).items())\n    else:\n      items.append((new_key, v))\n  return dict(items)\n\n\nclass SummaryWriter(object):\n  \"\"\"Saves data in event and summary protos for tensorboard.\"\"\"\n\n  def __init__(self, log_dir):\n    \"\"\"Create a new SummaryWriter.\n\n    Args:\n      log_dir: path to record tfevents files in.\n    \"\"\"\n    log_dir = os.fspath(log_dir)\n\n    # If needed, create log_dir directory as well as missing parent directories.\n    if not tf.io.gfile.isdir(log_dir):\n      tf.io.gfile.makedirs(log_dir)\n\n    self._event_writer = tf.summary.create_file_writer(log_dir, 10, 120, None)\n    self._closed = False\n\n  def close(self):\n    \"\"\"Close SummaryWriter. Final!\"\"\"\n    if not self._closed:\n      self._event_writer.close()\n      self._closed = True\n      del self._event_writer\n\n  def flush(self):\n    self._event_writer.flush()\n\n  def scalar(self, tag, value, step):\n    \"\"\"Saves scalar value.\n\n    Args:\n      tag: str: label for this data\n      value: int/float: number to log\n      step: int: training step\n    \"\"\"\n    value = float(np.array(value))\n    with self._event_writer.as_default():\n      tf.summary.scalar(name=tag, data=value, step=step)\n\n  def image(self, tag, image, step, max_outputs=3):\n    \"\"\"Saves RGB image summary from np.ndarray [H,W], [H,W,1], or [H,W,3].\n\n    Args:\n      tag: str: label for this data\n      image: ndarray: [H,W], [H,W,1], [H,W,3], [K,H,W], [K,H,W,1], [K,H,W,3]\n        Save image in greyscale or colors.\n        Pixel values could be either uint8 or float.\n        Floating point values should be in range [0, 1).\n      step: int: training step\n      max_outputs: At most this many images will be emitted at each step.\n        Defaults to 3.\n    \"\"\"\n    image = np.array(image)\n    # tf.summary.image expects image to have shape [k, h, w, c] where,\n    # k = number of samples, h = height, w = width, c = number of channels.\n    if len(np.shape(image)) == 2:\n      image = image[np.newaxis, :, :, np.newaxis]\n    elif len(np.shape(image)) == 3:\n      # this could be either [k, h, w] or [h, w, c]\n      if np.shape(image)[-1] in (1, 3):\n        image = image[np.newaxis, :, :, :]\n      else:\n        image = image[:, :, :, np.newaxis]\n    if np.shape(image)[-1] == 1:\n      image = np.repeat(image, 3, axis=-1)\n\n    # Convert to tensor value as tf.summary.image expects data to be a tensor.\n    image = tf.convert_to_tensor(image)\n    with self._event_writer.as_default():\n      tf.summary.image(name=tag, data=image, step=step, max_outputs=max_outputs)\n\n  def audio(self, tag, audiodata, step, sample_rate=44100, max_outputs=3):\n    \"\"\"Saves audio as wave.\n\n    NB: single channel only right now.\n\n    Args:\n      tag: str: label for this data\n      audiodata: ndarray [Nsamples, Nframes, Nchannels]: audio data to\n        be saved as wave. The data will be clipped to [-1.0, 1.0].\n      step: int: training step\n      sample_rate: sample rate of passed in audio buffer\n      max_outputs: At most this many audio clips will be emitted at each\n        step. Defaults to 3.\n    \"\"\"\n    # tf.summary.audio expects the audio data to have floating values in\n    # [-1.0, 1.0].\n    audiodata = np.clip(np.array(audiodata), -1, 1)\n\n    # Convert to tensor value as tf.summary.audio expects data to be a tensor.\n    audio = tf.convert_to_tensor(audiodata, dtype=tf.float32)\n    with self._event_writer.as_default():\n      tf.summary.audio(\n          name=tag, data=audio, sample_rate=sample_rate, step=step,\n          max_outputs=max_outputs, encoding='wav')\n\n  def histogram(self, tag, values, step, bins=None):\n    \"\"\"Saves histogram of values.\n\n    Args:\n      tag: str: label for this data\n      values: ndarray: will be flattened by this routine\n      step: int: training step\n      bins: number of bins in histogram\n    \"\"\"\n    values = np.array(values)\n    values = np.reshape(values, -1)\n    with self._event_writer.as_default():\n      tf.summary.histogram(name=tag, data=values, step=step, buckets=bins)\n\n  def text(self, tag, textdata, step):\n    \"\"\"Saves a text summary.\n\n    Args:\n      tag: str: label for this data\n      textdata: string\n      step: int: training step\n    Note: markdown formatting is rendered by tensorboard.\n    \"\"\"\n    if not isinstance(textdata, (str, bytes)):\n      raise ValueError('`textdata` should be of the type `str` or `bytes`.')\n    with self._event_writer.as_default():\n      tf.summary.text(name=tag, data=tf.constant(textdata), step=step)\n\n  def write(self, tag, tensor, step, metadata=None):\n    \"\"\"Saves a arbitrary tensor summary.\n\n    Useful when working with custom plugins or constructing a summary directly.\n\n    Args:\n      tag: str: label for this data\n      tensor: ndarray: tensor data to save.\n      step: int: training step\n      metadata: Optional SummaryMetadata, as a proto or serialized bytes.\n    Note: markdown formatting is rendered by tensorboard.\n    \"\"\"\n    with self._event_writer.as_default():\n      tf.summary.write(\n          tag=tag,\n          tensor=tensor,\n          step=step,\n          metadata=metadata)\n\n  def hparams(self, hparams):\n    \"\"\"Saves hyper parameters.\n\n    Args:\n      hparams: Flat mapping from hyper parameter name to value.\n    \"\"\"\n\n    with self._event_writer.as_default():\n      hparams_api.hparams(hparams=_flatten_dict(hparams))\n", "path": "flax/metrics/tensorboard.py"}], "after_files": [{"content": "# Copyright 2022 The Flax Authors.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Write Summaries from JAX for use with Tensorboard.\n\"\"\"\n\nimport os\n\n# pylint: disable=g-import-not-at-top\nimport numpy as np\n\nimport tensorflow.compat.v2 as tf\nfrom tensorboard.plugins.hparams import api as hparams_api\n\n\ndef _flatten_dict(input_dict, parent_key='', sep='.'):\n  \"\"\"Flattens and simplifies dict such that it can be used by hparams.\n\n  Args:\n    input_dict: Input dict, e.g., from ConfigDict.\n    parent_key: String used in recursion.\n    sep: String used to separate parent and child keys.\n\n  Returns:\n   Flattened dict.\n  \"\"\"\n  items = []\n  for k, v in input_dict.items():\n    new_key = parent_key + sep + k if parent_key else k\n\n    # Valid types according to https://github.com/tensorflow/tensorboard/blob/1204566da5437af55109f7a4af18f9f8b7c4f864/tensorboard/plugins/hparams/summary_v2.py\n    valid_types = (bool, int, float, str, np.bool_, np.integer, np.floating, np.character)\n\n    if isinstance(v, dict):\n      # Recursively flatten the dict.\n      items.extend(_flatten_dict(v, new_key, sep=sep).items())\n      continue\n    elif not isinstance(v, valid_types):\n      # Cast any incompatible values as strings such that they can be handled by hparams\n      v = str(v)\n    items.append((new_key, v))\n  return dict(items)\n\n\nclass SummaryWriter(object):\n  \"\"\"Saves data in event and summary protos for tensorboard.\"\"\"\n\n  def __init__(self, log_dir):\n    \"\"\"Create a new SummaryWriter.\n\n    Args:\n      log_dir: path to record tfevents files in.\n    \"\"\"\n    log_dir = os.fspath(log_dir)\n\n    # If needed, create log_dir directory as well as missing parent directories.\n    if not tf.io.gfile.isdir(log_dir):\n      tf.io.gfile.makedirs(log_dir)\n\n    self._event_writer = tf.summary.create_file_writer(log_dir, 10, 120, None)\n    self._closed = False\n\n  def close(self):\n    \"\"\"Close SummaryWriter. Final!\"\"\"\n    if not self._closed:\n      self._event_writer.close()\n      self._closed = True\n      del self._event_writer\n\n  def flush(self):\n    self._event_writer.flush()\n\n  def scalar(self, tag, value, step):\n    \"\"\"Saves scalar value.\n\n    Args:\n      tag: str: label for this data\n      value: int/float: number to log\n      step: int: training step\n    \"\"\"\n    value = float(np.array(value))\n    with self._event_writer.as_default():\n      tf.summary.scalar(name=tag, data=value, step=step)\n\n  def image(self, tag, image, step, max_outputs=3):\n    \"\"\"Saves RGB image summary from np.ndarray [H,W], [H,W,1], or [H,W,3].\n\n    Args:\n      tag: str: label for this data\n      image: ndarray: [H,W], [H,W,1], [H,W,3], [K,H,W], [K,H,W,1], [K,H,W,3]\n        Save image in greyscale or colors.\n        Pixel values could be either uint8 or float.\n        Floating point values should be in range [0, 1).\n      step: int: training step\n      max_outputs: At most this many images will be emitted at each step.\n        Defaults to 3.\n    \"\"\"\n    image = np.array(image)\n    # tf.summary.image expects image to have shape [k, h, w, c] where,\n    # k = number of samples, h = height, w = width, c = number of channels.\n    if len(np.shape(image)) == 2:\n      image = image[np.newaxis, :, :, np.newaxis]\n    elif len(np.shape(image)) == 3:\n      # this could be either [k, h, w] or [h, w, c]\n      if np.shape(image)[-1] in (1, 3):\n        image = image[np.newaxis, :, :, :]\n      else:\n        image = image[:, :, :, np.newaxis]\n    if np.shape(image)[-1] == 1:\n      image = np.repeat(image, 3, axis=-1)\n\n    # Convert to tensor value as tf.summary.image expects data to be a tensor.\n    image = tf.convert_to_tensor(image)\n    with self._event_writer.as_default():\n      tf.summary.image(name=tag, data=image, step=step, max_outputs=max_outputs)\n\n  def audio(self, tag, audiodata, step, sample_rate=44100, max_outputs=3):\n    \"\"\"Saves audio as wave.\n\n    NB: single channel only right now.\n\n    Args:\n      tag: str: label for this data\n      audiodata: ndarray [Nsamples, Nframes, Nchannels]: audio data to\n        be saved as wave. The data will be clipped to [-1.0, 1.0].\n      step: int: training step\n      sample_rate: sample rate of passed in audio buffer\n      max_outputs: At most this many audio clips will be emitted at each\n        step. Defaults to 3.\n    \"\"\"\n    # tf.summary.audio expects the audio data to have floating values in\n    # [-1.0, 1.0].\n    audiodata = np.clip(np.array(audiodata), -1, 1)\n\n    # Convert to tensor value as tf.summary.audio expects data to be a tensor.\n    audio = tf.convert_to_tensor(audiodata, dtype=tf.float32)\n    with self._event_writer.as_default():\n      tf.summary.audio(\n          name=tag, data=audio, sample_rate=sample_rate, step=step,\n          max_outputs=max_outputs, encoding='wav')\n\n  def histogram(self, tag, values, step, bins=None):\n    \"\"\"Saves histogram of values.\n\n    Args:\n      tag: str: label for this data\n      values: ndarray: will be flattened by this routine\n      step: int: training step\n      bins: number of bins in histogram\n    \"\"\"\n    values = np.array(values)\n    values = np.reshape(values, -1)\n    with self._event_writer.as_default():\n      tf.summary.histogram(name=tag, data=values, step=step, buckets=bins)\n\n  def text(self, tag, textdata, step):\n    \"\"\"Saves a text summary.\n\n    Args:\n      tag: str: label for this data\n      textdata: string\n      step: int: training step\n    Note: markdown formatting is rendered by tensorboard.\n    \"\"\"\n    if not isinstance(textdata, (str, bytes)):\n      raise ValueError('`textdata` should be of the type `str` or `bytes`.')\n    with self._event_writer.as_default():\n      tf.summary.text(name=tag, data=tf.constant(textdata), step=step)\n\n  def write(self, tag, tensor, step, metadata=None):\n    \"\"\"Saves a arbitrary tensor summary.\n\n    Useful when working with custom plugins or constructing a summary directly.\n\n    Args:\n      tag: str: label for this data\n      tensor: ndarray: tensor data to save.\n      step: int: training step\n      metadata: Optional SummaryMetadata, as a proto or serialized bytes.\n    Note: markdown formatting is rendered by tensorboard.\n    \"\"\"\n    with self._event_writer.as_default():\n      tf.summary.write(\n          tag=tag,\n          tensor=tensor,\n          step=step,\n          metadata=metadata)\n\n  def hparams(self, hparams):\n    \"\"\"Saves hyper parameters.\n\n    Args:\n      hparams: Flat mapping from hyper parameter name to value.\n    \"\"\"\n\n    with self._event_writer.as_default():\n      hparams_api.hparams(hparams=_flatten_dict(hparams))\n", "path": "flax/metrics/tensorboard.py"}]}