code: def sum_reward_weighted_observations(
      r: types.Tensor, x: types.Tensor
  ) -> types.Tensor:
      """Calculates an update used by some Bandit algorithms.
  
      Given an observation `x` and corresponding reward `r`, the weigthed
      observations vector (denoted `b` here) should be updated as `b = b + r * x`.
      This function calculates the sum of weighted rewards for batched
      observations `x`.
  
      Args:
        r: a `Tensor` of shape [`batch_size`]. This is the rewards of the batched
          observations.
        x: a `Tensor` of shape [`batch_size`, `context_dim`]. This is the matrix
          with the (batched) observations.
  
      Returns:
        The update that needs to be added to `b`. Has the same shape as `b`.
        If the observation matrix `x` is empty, a zero vector is returned.
      """
      batch_size = tf.shape(x)[0]
  
      return tf.reduce_sum(tf.reshape(r, [batch_size, 1]) * x, axis=0)
  
file_path: agents/tf_agents/bandits/agents/utils.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Tuple
  import gin
  import numpy
  import tensorflow
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import utils
  from tf_agents.typing import types
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 1 for calling a global function with example usage: sum_reward_weighted_observations(r, x) and returns: tf