code: def scalarize_objectives(
      objectives_tensor: tf.Tensor,
      scalarizer: multi_objective_scalarizer.Scalarizer,
  ):
      """Scalarize a rank-3 objectives tensor into a rank-2 tensor.
  
      Scalarize an objective values tensor shaped as
      [batch_size, num_of_objectives, num_of_actions] along the second dimension
      into a rank-2 tensor shaped as [batch_size, num_of_actions]
  
      Args:
        objectives_tensor: An objectives tensor to be scalarized.
        scalarizer: A
          `tf_agents.bandits.multi_objective.multi_objective_scalarizer.Scalarizer`
          object that implements scalarization of multiple objectives into a single
          scalar reward.
  
      Returns:
        A rank-2 tensor of scalarized rewards shaped as
        [batch_size, num_of_actions].
  
      Raises:
        ValueError: If `objectives_tensor` is not rank-3.
      """
      if objectives_tensor.shape.rank != 3:
          raise ValueError(
              "The objectives_tensor should be rank-3, but is rank-{}".format(
                  objectives_tensor.shape.rank
              )
          )
      return tf.transpose(
          tf.map_fn(
              scalarizer,
              tf.transpose(objectives_tensor, perm=[2, 0, 1]),
          )
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import copy
  from typing import List, Optional, Sequence, Text, Tuple
  import gin
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.multi_objective import multi_objective_scalarizer
  from tf_agents.bandits.networks import heteroscedastic_q_network
  from tf_agents.bandits.policies import constraints
  from tf_agents.bandits.specs import utils
  from tf_agents.networks.network import Network
  from tf_agents.policies import tf_policy
  from tf_agents.policies import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import policy_step
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
label_desc: Use image node_id 1 for calling a global function with example usage: scalarize_objectives(objectives_tensor, scalarizer) and returns: tf