code: def testBUpdate(self, batch_size, context_dim):
      b_array = np.array(range(context_dim))
      r_array = np.array(range(batch_size)).reshape((batch_size, 1))
      x_array = np.array(range(batch_size * context_dim)).reshape(
          (batch_size, context_dim)
      )
      rx = r_array * x_array
      expected_b_updated_array = b_array + np.sum(rx, axis=0)
  
      b = tf.constant(b_array, dtype=tf.float32, shape=[context_dim])
      r = tf.constant(r_array, dtype=tf.float32, shape=[batch_size])
      x = tf.constant(
          x_array, dtype=tf.float32, shape=[batch_size, context_dim]
      )
      b_update = utils.sum_reward_weighted_observations(r, x)
      self.assertAllClose(
          expected_b_updated_array, self.evaluate(b + b_update)
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.agents import utils
  from tf_agents.specs import tensor_spec
  from tf_agents.trajectories import time_step
  from tf_agents.trajectories import trajectory
label_desc: Use image node_id 1 for calling the UtilsTest obj's underlying member method code with example usage: obj.testBUpdate(batch_size, context_dim) without return types