code: def _dtensor_variable_summary(weights):
      """Group and calculate DTensor based weights memory size.
  
      Since DTensor weights can be sharded across multiple device, the result
      will be grouped by the layout/sharding spec for the variables, so that
      the accurate per-device memory size can be calculated.
  
      Args:
          weights: An iterable contains the weights to compute weight size.
  
      Returns:
          total_weight_count, total_memory_size and per_sharing_spec_result which
          is a dict with normalized layout spec as key and tuple of weight count
          and weight size as value.
      """
      unique_weights = {id(w): w for w in weights}.values()
      total_weight_count = 0
      total_memory_size = 0
      per_sharing_spec_result = {}
      for w in unique_weights:
          # Ignore TrackableWeightHandlers, which will not have a shape defined.
          if not hasattr(w, "shape"):
              continue
          if not isinstance(w, tf.experimental.dtensor.DVariable):
              continue
          layout = w.layout
          # Remove all the duplication axis, and sort the column name.
          # 1D replicated and 2D replicated variable will still be fully
          # replicated, and [batch, model] sharding will have same memory
          # footprint as the [model, batch] layout.
          reduced_sharding_spec = list(
              sorted(set(layout.sharding_specs))
          )
          if tf.experimental.dtensor.UNSHARDED in reduced_sharding_spec:
              reduced_sharding_spec.remove(
                  tf.experimental.dtensor.UNSHARDED
              )
          reduced_sharding_spec = tuple(
              reduced_sharding_spec
          )  # For dict key
          weight_count, memory_size = per_sharing_spec_result.get(
              reduced_sharding_spec, (0, 0)
          )
          reduced_weight_shape = np.prod(w.shape.as_list())
          per_param_size = w.dtype.size
          weight_count += reduced_weight_shape
          memory_size += reduced_weight_shape * per_param_size
          per_sharing_spec_result[reduced_sharding_spec] = (
              weight_count,
              memory_size,
          )
          total_weight_count += reduced_weight_shape
          total_memory_size += reduced_weight_shape * per_param_size
      return (
          total_weight_count,
          total_memory_size,
          per_sharing_spec_result,
      )
  
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 5 for calling a global function with example usage: _dtensor_variable_summary(weights) and returns: total_weight_count,  total_memory_size,  per_sharing_spec_result