def get_chart_table_args(stats_table, device_type):
  """Creates gviz DataTable object from a TensorFlow stats table.

  Args:
    stats_table: A tf_stats_pb2.TfStatsTable.
    device_type: A string representing device type of this profile.

  Returns:
    Returns a gviz_api.DataTable
  """

  # Create schema
  table_description = [
      ("rank", "number", "Rank"),
      ("host_or_device", "string", "Host/device"),
      ("type", "string", "Type"),
      ("operation", "string", "Operation"),
      ("occurrences", "number", "#Occurrences"),
      ("total_time", "number", "Total time (us)"),
      ("avg_time", "number", "Avg. time (us)"),
      ("total_self_time", "number", "Total self-time (us)"),
      ("avg_self_time", "number", "Avg. self-time (us)"),
      ("device_total_self_time_percent", "number",
       "Total self-time on Device (%)"),
      ("device_cumulative_total_self_time_percent", "number",
       "Cumulative total-self time on Device (%)"),
      ("host_total_self_time_percent", "number", "Total self-time on Host (%)"),
      ("host_cumulative_total_self_time_percent", "number",
       "Cumulative total-self time on Host (%)"),
  ]

  # Add GPU TensorCore utilization only when hardware device type is GPU, for
  # other devices, the GPU TensorCore utilization is always 0 so there is no
  # need to display this column.
  if device_type == "GPU":
    table_description.append(
        ("gpu_tensorcore_utilization", "number", "GPU TensorCore utilization"))

  data = []
  for record in stats_table.tf_stats_record:
    row = [
        record.rank,
        record.host_or_device,
        record.op_type,
        record.op_name,
        record.occurrences,
        record.total_time_in_us,
        record.avg_time_in_us,
        record.total_self_time_in_us,
        record.avg_self_time_in_us,
        record.device_total_self_time_as_fraction,
        record.device_cumulative_total_self_time_as_fraction,
        record.host_total_self_time_as_fraction,
        record.host_cumulative_total_self_time_as_fraction,
    ]

    if device_type == "GPU":
      row.append(record.gpu_tensorcore_utilization)

    data.append(row)

  return (table_description, data, [])