def _replace_nan_with_none(
    plot_data,
    plot_keys):
  """Replaces all instances of nan with None in plot data.

  This is necessary for Colab integration where we serializes the data into json
  string as NaN is not supported by json standard. Turning nan into None will
  make the value null once parsed. The visualization already handles falsy
  values by setting them to zero.

  Args:
    plot_data: The original plot data
    plot_keys: A dictionary containing field names of plot data.

  Returns:
    Transformed plot data where all nan has been replaced with None.
  """
  output_metrics = {}
  for plot_type in plot_keys:
    metric_name = plot_keys[plot_type]['metricName']
    if metric_name in plot_data:
      data_series_name = plot_keys[plot_type]['dataSeries']
      if data_series_name in plot_data[metric_name]:
        data_series = plot_data[metric_name][data_series_name]
        outputs = []
        for entry in data_series:
          output = {}
          for key in entry:
            value = entry[key]
            # When converting protocol buffer into dict, float value nan is
            # automatically converted into the string 'NaN'.
            output[key] = None if value == 'NaN' else value
          outputs.append(output)
        output_metrics[metric_name] = {data_series_name: outputs}

  return output_metrics