class_docstr: An agent implementing the LinUCB algorithm on top of a neural network.
docstr: Computes the loss using LinUCB distributively.
  
  Args:
    observation: A batch of observations.
    action: A batch of actions.
    reward: A batch of rewards.
    weights: unused weights.
    training: Whether the loss is being used to train.
  
  Returns:
    loss: A `LossInfo` containing the loss for the training step.
docstr_tok: ["Computes","the","loss","using","LinUCB","distributively",".","Args",":","observation",":","A","batch","of","observations",".","action",":","A","batch","of","actions",".","reward",":","A","batch","of","rewards",".","weights",":","unused","weights",".","training",":","Whether","the","loss","is","being","used","to","train",".","Returns",":","loss",":","A","`","LossInfo","`","containing","the","loss","for","the","training","step","."]
code: def compute_loss_using_linucb_distributed(
      self,
      observation: types.NestedTensor,
      action: types.Tensor,
      reward: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> tf_agent.LossInfo:
      """Computes the loss using LinUCB distributively.
  
      Args:
        observation: A batch of observations.
        action: A batch of actions.
        reward: A batch of rewards.
        weights: unused weights.
        training: Whether the loss is being used to train.
  
      Returns:
        loss: A `LossInfo` containing the loss for the training step.
      """
      del weights  # unused
  
      # The network is trained now. Update the covariance matrix.
      encoded_observation, _ = self._encoding_network(
          observation, training=training
      )
      encoded_observation = tf.cast(
          encoded_observation, dtype=self._dtype
      )
      encoded_observation = tf.reshape(
          encoded_observation, shape=[-1, self._encoding_dim]
      )
  
      self._train_step_counter.assign_add(1)
  
      for k in range(self._num_models):
          diag_mask = tf.linalg.tensor_diag(
              tf.cast(tf.equal(action, k), self._dtype)
          )
          observations_for_arm = tf.matmul(
              diag_mask, encoded_observation
          )
          rewards_for_arm = tf.matmul(
              diag_mask, tf.reshape(reward, [-1, 1])
          )
  
          # Compute local updates for the matrix A and b of this arm.
          cov_matrix_local_udpate = tf.matmul(
              observations_for_arm,
              observations_for_arm,
              transpose_a=True,
          )
          data_vector_local_update = (
              bandit_utils.sum_reward_weighted_observations(
                  rewards_for_arm, observations_for_arm
              )
          )
  
          def _merge_fn(
              strategy,
              per_replica_cov_matrix_update,
              per_replica_data_vector_update,
          ):
              """Merge the per-replica-updates."""
              # Reduce the per-replica-updates using SUM.
              # pylint: disable=cell-var-from-loop
              updates_and_vars = [
                  (per_replica_cov_matrix_update, self.cov_matrix[k]),
                  (per_replica_data_vector_update, self.data_vector[k]),
              ]
  
              reduced_updates = strategy.extended.batch_reduce_to(
                  tf.distribute.ReduceOp.SUM, updates_and_vars
              )
  
              # Update the model variables.
              self.cov_matrix[k].assign(
                  self._gamma * self.cov_matrix[k] + reduced_updates[0]
              )
              self.data_vector[k].assign(
                  self._gamma * self.data_vector[k] + reduced_updates[1]
              )
  
          # Passes the local_updates to the _merge_fn() above that performs custom
          # computation on the per-replica values.
          # All replicas pause their execution until merge_call() is done and then,
          # execution is resumed.
          replica_context = tf.distribute.get_replica_context()
          replica_context.merge_call(
              _merge_fn,
              args=(cov_matrix_local_udpate, data_vector_local_update),
          )
  
      loss = -1.0 * tf.reduce_sum(reward)
      return tf_agent.LossInfo(loss=(loss), extra=())
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from typing import Optional, Sequence, Text
  import gin
  import tensorflow
  from tf_agents.agents import data_converter
  from tf_agents.agents import tf_agent
  from tf_agents.bandits.agents import linear_bandit_agent
  from tf_agents.bandits.agents import utils
  from tf_agents.bandits.policies import neural_linucb_policy
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import utils
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import eager_utils
label_desc: Use image node_id 15 for calling the NeuralLinUCBAgent obj's underlying member method code with example usage: obj.compute_loss_using_linucb_distributed(observation, action, reward, weights, training) and returns: tf_agent