class_docstr: A simple neural network with dense layers and sigmoid non-linearity.
  
  The network consists of `len(hidden_layers) + 1` dense layers. The sizes of
  the hidden layers are specified by the user in `hidden_layers` and the
  network takes care of adding layers to match the input and output size.
  
  Attributes:
    weights: A list of 2-d float32 arrays containing the layer weights.
    biases: A list of 2-d float32 arrays containing the layer biases.
  
  Methods:
    forward: Can be used to perform a forward pass on a batch of
      flattened images. Output is returned as a batch of one-hot vectors of the
      classes.
    train: method performs a forward and backward pass and updates the
      weights and biases.
    evaluate: method can be used to evaluate the network on a batch of
      examples.
docstr: 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.
code: 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 :]
  
file_path: trax/trax/tf_numpy/examples/mnist/model.py
imports: from numpy import float32
  from numpy import int32
  import tensorflow.compat.v2
  from trax.tf_numpy import numpy
label_desc: Use image node_id 3 for calling the Model obj's underlying member method code with example usage: obj.train(x, y, learning_rate) and returns: unknown