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))
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.
    
trax/tf_numpy/examples/mnist/model.py