class_docstr: Base class for coding experiment.
class_docstr_tok: ["Base","class","for","coding","experiment","."]
code: def train_step(self, x):
      if hasattr(self, "alpha"):
          self.alpha = self.force_alpha
      with tf.GradientTape() as tape:
          rates, distortions = self.train_losses(x)
          losses = rates + self.lmbda * distortions
          loss = tf.math.reduce_mean(losses)
      variables = self.trainable_variables
      gradients = tape.gradient(loss, variables)
      self.optimizer.apply_gradients(zip(gradients, variables))
      self.loss.update_state(losses)
      self.rate.update_state(rates)
      self.distortion.update_state(distortions)
      energy = []
      size = []
      for grad in gradients:
          if grad is None:
              continue
          energy.append(
              tf.reduce_sum(tf.square(tf.cast(grad, tf.float64)))
          )
          size.append(tf.cast(tf.size(grad), tf.float64))
      self.grad_rms.update_state(
          tf.sqrt(tf.add_n(energy) / tf.add_n(size))
      )
      return {
          m.name: m.result()
          for m in [
              self.loss,
              self.rate,
              self.distortion,
              self.grad_rms,
          ]
      }
  
file_path: compression/models/toy_sources/compression_model.py
imports: import abc
  import matplotlib.pyplot
  import numpy
  import tensorflow
label_desc: Use image node_id 7 for calling the CompressionModel obj's underlying member method code with example usage: obj.train_step(x) and returns: unknown