docstr: Same as above but with distributed LinUCB updates.
docstr_tok: ["Same","as","above","but","with","distributed","LinUCB","updates","."]
code: def testNeuralLinUCBUpdateDistributed(
      self, batch_size=1, context_dim=10
  ):
      """Same as above but with distributed LinUCB updates."""
  
      # Construct a `Trajectory` for the given action, observation, reward.
      num_actions = 5
      initial_step, final_step = _get_initial_and_final_steps(
          batch_size, context_dim
      )
      action = np.random.randint(
          num_actions, size=batch_size, dtype=np.int32
      )
      action_step = _get_action_step(action)
      experience = _get_experience(
          initial_step, action_step, final_step
      )
  
      # Construct an agent and perform the update.
      observation_spec = tensor_spec.TensorSpec(
          [context_dim], tf.float32
      )
      time_step_spec = time_step.time_step_spec(observation_spec)
      action_spec = tensor_spec.BoundedTensorSpec(
          dtype=tf.int32, shape=(), minimum=0, maximum=num_actions - 1
      )
      encoder = DummyNet(observation_spec)
      encoding_dim = 10
      agent = neural_linucb_agent.NeuralLinUCBAgent(
          time_step_spec=time_step_spec,
          action_spec=action_spec,
          encoding_network=encoder,
          encoding_network_num_train_steps=0,
          encoding_dim=encoding_dim,
          optimizer=tf.compat.v1.train.AdamOptimizer(
              learning_rate=1e-2
          ),
      )
  
      self.evaluate(agent.initialize())
      self.evaluate(tf.compat.v1.global_variables_initializer())
      # Call the distributed LinUCB training instead of agent.train().
      train_fn = common.function_in_tf1()(
          agent.compute_loss_using_linucb_distributed
      )
      reward = tf.cast(experience.reward, agent._dtype)
      loss_info = train_fn(
          experience.observation, action, reward, weights=None
      )
      self.evaluate(loss_info)
      final_a = self.evaluate(agent.cov_matrix)
      final_b = self.evaluate(agent.data_vector)
  
      # Compute the expected updated estimates.
      observations_list = tf.dynamic_partition(
          data=tf.reshape(
              experience.observation, [batch_size, context_dim]
          ),
          partitions=tf.convert_to_tensor(action),
          num_partitions=num_actions,
      )
      rewards_list = tf.dynamic_partition(
          data=tf.reshape(experience.reward, [batch_size]),
          partitions=tf.convert_to_tensor(action),
          num_partitions=num_actions,
      )
      expected_a_updated_list = []
      expected_b_updated_list = []
      for _, (observations_for_arm, rewards_for_arm) in enumerate(
          zip(observations_list, rewards_list)
      ):
          encoded_observations_for_arm, _ = encoder(
              observations_for_arm
          )
  
          num_samples_for_arm_current = tf.cast(
              tf.shape(rewards_for_arm)[0], tf.float32
          )
          num_samples_for_arm_total = num_samples_for_arm_current
  
          # pylint: disable=cell-var-from-loop
          def true_fn():
              a_new = tf.matmul(
                  encoded_observations_for_arm,
                  encoded_observations_for_arm,
                  transpose_a=True,
              )
              b_new = bandit_utils.sum_reward_weighted_observations(
                  rewards_for_arm, encoded_observations_for_arm
              )
              return a_new, b_new
  
          def false_fn():
              return (
                  tf.zeros(
                      [encoding_dim, encoding_dim], dtype=tf.float32
                  ),
                  tf.zeros([encoding_dim], dtype=tf.float32),
              )
  
          a_new, b_new = tf.cond(
              tf.squeeze(num_samples_for_arm_total) > 0,
              true_fn,
              false_fn,
          )
  
          expected_a_updated_list.append(self.evaluate(a_new))
          expected_b_updated_list.append(self.evaluate(b_new))
  
      # Check that the actual updated estimates match the expectations.
      self.assertAllClose(expected_a_updated_list, final_a)
      self.assertAllClose(expected_b_updated_list, final_b)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import os
  from absl.testing import parameterized
  import numpy
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.agents import neural_linucb_agent
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.drivers import driver_utils
  from tf_agents.bandits.networks import global_and_arm_feature_network
  from tf_agents.bandits.specs import utils
  from tf_agents.networks import network
  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.utils import common
label_desc: Use image node_id 5 for calling the NeuralLinUCBAgentTest obj's underlying member method code with example usage: obj.testNeuralLinUCBUpdateDistributed(batch_size, context_dim) and returns: a_new,  b_new, tf,  tf