code: def _print_dtensor_variable_summary(model, print_fn, line_length):
      """Prints dtensor variables summary."""
      if (
          getattr(model, "_layout_map", None) is not None
      ):  # pylint: disable=protected-access
          mesh = (
              model._layout_map.get_default_mesh()
          )  # pylint: disable=protected-access
      elif hasattr(model, "distribute_strategy") and hasattr(
          model.distribute_strategy, "_mesh"
      ):
          mesh = (
              model.distribute_strategy._mesh
          )  # pylint: disable=protected-access
      else:
          # Not running with DTensor
          mesh = None
      if mesh:
          (
              total_weight_count,
              total_memory_size,
              per_sharing_spec_result,
          ) = _dtensor_variable_summary(model.weights)
          total_per_device_memory_size = 0
          for sharding_spec in sorted(
              per_sharing_spec_result.keys()
          ):  # pylint: disable=g-builtin-op
              count, memory_size = per_sharing_spec_result[
                  sharding_spec
              ]
              if (
                  len(sharding_spec) == 0
              ):  # pylint: disable=g-explicit-length-test
                  print_fn(
                      f"{count} / {total_weight_count} params "
                      f"({_readable_memory_size(memory_size)}) "
                      "are fully replicated"
                  )
                  per_device_size = memory_size
              else:
                  sharding_factor = np.prod(
                      [mesh.dim_size(s) for s in sharding_spec]
                  )
                  per_device_size = memory_size / sharding_factor
                  print_fn(
                      f"{count} / {total_weight_count} params "
                      f"({_readable_memory_size(memory_size)}) are sharded based "
                      f"on spec '{sharding_spec}' and across {sharding_factor} "
                      "devices."
                  )
              total_per_device_memory_size += per_device_size
          print_fn(
              "Overall per device memory usage: "
              f"{_readable_memory_size(total_per_device_memory_size)}"
          )
          print_fn(
              "Overall sharding factor: {:.2f}".format(
                  total_memory_size / total_per_device_memory_size
              )
          )
          print_fn("_" * line_length)
  
file_path: agents/tf_agents/networks/layer_utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import re
  import numpy
  import tensorflow
label_desc: Use image node_id 6 for calling a global function with example usage: _print_dtensor_variable_summary(model, print_fn, line_length) without return types