code: def batchwise_scalar_to_metric_with_counts(
      metric, weights, devicewise_counts
  ):
      norm_metric = train_utils.psum_metric_normalizer(
          (metric * weights, devicewise_counts)
      )
      return norm_metric[0], devicewise_counts
  
imports: import functools
  import jax.numpy
  from scenic.projects.adversarialtraining.attacks import train_utils
label_desc: Use image node_id 1 for calling a global function with example usage: batchwise_scalar_to_metric_with_counts(metric, weights, devicewise_counts) and returns: norm_metric,  devicewise_counts