//  [python]
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)