class_docstr: Implementation of control theory augment.
class_docstr_tok: ["Implementation","of","control","theory","augment","."]
docstr: Update augmenter state to classification of probe images.
docstr_tok: ["Update","augmenter","state","to","classification","of","probe","images","."]
code: def update_rates(self, tensor_dict, probe_probs):
      """Update augmenter state to classification of probe images."""
      # shape of probe_probs is (batch_size, num_classes)
      op_idx = tensor_dict[
          "probe_op_indices"
      ]  # shape=(batch_size, num_layers)
      op_arg = tensor_dict[
          "probe_op_args"
      ]  # shape=(batch_size, num_layers)
      label = tf.expand_dims(
          tensor_dict["label"], 1
      )  # shape=(batch_size, 1)
  
      # Compute proximity metric as softmax(model(probe_image))[correct_label]
      # Tile proximity, so its shape will be (batch_size, num_layers)
      proximity = tf.gather(probe_probs, label, axis=1, batch_dims=1)
      proximity = tf.tile(proximity, [1, self.n])
      # Quantize op_arg to obtain levels of the ops.
      # NOTE: computed level should be always less than num_levels,
      #       nevertherless use minimum operation to enforce the range.
      level_idx = tf.cast(op_arg * self.quantum, tf.int32)
      level_idx = tf.minimum(level_idx, self.quantum)
  
      # Update rates.
      # For each (op_index, level_index, proximity) in the list of selected ops
      # update rate using following formula:
      #   rate[op_idx, level_idx] = rate[op_idx, level_idx] * decay
      #                             + proximity * (1 - decay)
      # which is equivalent to:
      #   alpha = 1 - decay
      #   rate[op_idx, level_idx] += (proximity - rate[op_idx, level_idx]) * alpha
      #
      # So update is performed using assign_add operation. If several updates
      # correpond to the same (op_idx, level_idx) then they are averaged.
      op_level_idx = tf.concat(
          [tf.reshape(op_idx, [-1, 1]), tf.reshape(level_idx, [-1, 1])],
          axis=1,
      )
      flat_proximity = tf.reshape(proximity, [-1])
      sparse_update = (
          flat_proximity - tf.gather_nd(self.rates, op_level_idx)
      ) * self.alpha
      # Dense matrix with updates is computed in dense_update_numerator.
      # tf.scatter_nd adds up all updates which correspond to the same index,
      # however we need to compute mean. Thus we compute number of
      # updates corresponding to each index and divide by this number.
      dense_update_numerator = tf.scatter_nd(
          op_level_idx, sparse_update, shape=self.state_shape
      )
      dense_update_denominator = tf.scatter_nd(
          op_level_idx,
          tf.ones_like(sparse_update),
          shape=self.state_shape,
      )
      dense_update_denominator = tf.maximum(
          dense_update_denominator, 1.0
      )
      self.rates.assign_add(
          dense_update_numerator / dense_update_denominator
      )
  
      # Convert rates to log probabilities
      probs = tf.maximum(self.rates, self.eps)
      probs = probs / tf.reduce_max(probs, axis=1, keepdims=True)
      probs = tf.where(
          probs < self.confidence_threshold, tf.zeros_like(probs), probs
      )
      probs = probs + self.eps
      probs = probs / tf.reduce_sum(probs, axis=1, keepdims=True)
      self.probs.assign(probs)
  
imports: from typing import Optional
  import tensorflow
  from .randaugment import augment_ops
label_desc: Use image node_id 2 for calling the CTAugment obj's underlying member method code with example usage: obj.update_rates(tensor_dict, probe_probs) without return types