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: Initializes the neural network.
  
  Args:
    hidden_layers: List of ints specifying the sizes of hidden layers. Could
      be empty.
    input_size: Length of the input array. The network receives the input
      image as a flattened 1-d array. Defaults to 784(28*28), the default
      image size for MNIST.
    num_classes: The number of output classes. Defaults to 10.
code: def __init__(self, hidden_layers, input_size=784, num_classes=10):
      """Initializes the neural network.
  
      Args:
        hidden_layers: List of ints specifying the sizes of hidden layers. Could
          be empty.
        input_size: Length of the input array. The network receives the input
          image as a flattened 1-d array. Defaults to 784(28*28), the default
          image size for MNIST.
        num_classes: The number of output classes. Defaults to 10.
      """
      hidden_layers = [input_size] + hidden_layers + [num_classes]
      self.weights = []
      self.biases = []
      for i in range(len(hidden_layers) - 1):
          # TODO(srbs): This is manually cast to float32 to avoid the cast in
          # np.dot since backprop fails for tf.cast op.
          self.weights.append(
              np.array(
                  np.random.randn(
                      hidden_layers[i + 1], hidden_layers[i]
                  ),
                  copy=False,
                  dtype=float32,
              )
          )
          self.biases.append(
              np.array(
                  np.random.randn(hidden_layers[i + 1]),
                  copy=False,
                  dtype=float32,
              )
          )
  
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 1 to create a new Model object from inherited base classes: object with example: obj = Model(hidden_layers, input_size, num_classes)