def train(self, x, y, learning_rate=0.01):
    """Runs a single training pass.

    Args:
      x: 2-d array of size batch_size x image_size.
      y: 2-d array of size batch_size x num_classes in one-hot notation.
      learning_rate: The learning rate.
    """
    x = np.array(x, copy=False)
    y = np.array(y, copy=False)

    def mean_squared_error(x, y):
      diff = x - y
      return np.sum(diff * diff) / len(x)

    wb_tensors = self.weights + self.biases
    with tf.GradientTape() as g:
      g.watch(wb_tensors)
      loss = mean_squared_error(self.forward(x), y)
    gradients = g.gradient(loss, wb_tensors)
    gradients = [np.asarray(grad) for grad in gradients]

    new_weights_and_biases = []
    for v, dv in zip(self.weights + self.biases, gradients):
      new_weights_and_biases.append(v - learning_rate * dv)

    total_len = len(new_weights_and_biases)
    self.weights = new_weights_and_biases[:total_len // 2]
    self.biases = new_weights_and_biases[total_len // 2:]
Runs a single training pass.

    Args:
      x: 2-d array of size batch_size x image_size.
      y: 2-d array of size batch_size x num_classes in one-hot notation.
      learning_rate: The learning rate.
    
trax/tf_numpy/examples/mnist/model.py