class_docstr: Class for representing a trainable constraint using a neural network.
  
  This constraint class uses a neural network to compute the action feasibility.
  In this case, the loss function needs to be exposed for training the neural
  network weights, typically done by the agent that uses this constraint.
class_docstr_tok: ["Class","for","representing","a","trainable","constraint","using","a","neural","network",".","This","constraint","class","uses","a","neural","network","to","compute","the","action","feasibility",".","In","this","case",",","the","loss","function","needs","to","be","exposed","for","training","the","neural","network","weights",",","typically","done","by","the","agent","that","uses","this","constraint","."]
docstr: Computes loss for training the constraint network.
  
  Args:
    observations: A batch of observations.
    actions: A batch of actions.
    rewards: A batch of rewards.
    weights: Optional scalar or elementwise (per-batch-entry) importance
      weights.  The output batch loss will be scaled by these weights, and the
      final scalar loss is the mean of these values.
    training: Whether the loss is being used for training.
  
  Returns:
    loss: A `Tensor` containing the loss for the training step.
code: def compute_loss(
      self,
      observations: types.NestedTensor,
      actions: types.NestedTensor,
      rewards: types.Tensor,
      weights: Optional[types.Float] = None,
      training: bool = False,
  ) -> types.Tensor:
      """Computes loss for training the constraint network.
  
      Args:
        observations: A batch of observations.
        actions: A batch of actions.
        rewards: A batch of rewards.
        weights: Optional scalar or elementwise (per-batch-entry) importance
          weights.  The output batch loss will be scaled by these weights, and the
          final scalar loss is the mean of these values.
        training: Whether the loss is being used for training.
  
      Returns:
        loss: A `Tensor` containing the loss for the training step.
      """
      with tf.name_scope("constraint_loss"):
          sample_weights = weights if weights is not None else 1
          predicted_values, _ = self._constraint_network(
              observations, training=training
          )
          action_predicted_values = common.index_with_actions(
              predicted_values, tf.cast(actions, dtype=tf.int32)
          )
          # Reduction is done outside of the loss function because non-scalar
          # weights with unknown shapes may trigger shape validation that fails
          # XLA compilation.
          return tf.reduce_mean(
              tf.multiply(
                  self._error_loss_fn(
                      rewards,
                      action_predicted_values,
                      reduction=tf.compat.v1.losses.Reduction.NONE,
                  ),
                  sample_weights,
              )
          )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import abc
  import functools
  from typing import Callable, Iterable, Optional, Text
  import six
  import tensorflow
  import tensorflow_probability
  from tf_agents.bandits.policies import loss_utils
  from tf_agents.bandits.specs import utils
  from tf_agents.policies import utils
  from tf_agents.trajectories import time_step
  from tf_agents.typing import types
  from tf_agents.utils import common
  from tf_agents.utils import nest_utils
label_desc: Use image node_id 7 for calling the NeuralConstraint obj's underlying member method code with example usage: obj.compute_loss(observations, actions, rewards, weights, training) and returns: tf