def forward(self, x):
    """Performs the forward pass.

    Args:
      x: 2-d array of size batch_size x image_size.

    Returns:
      A 2-d array of size batch_size x num_classes.
    """

    def sigmoid(x):
      return 1.0 / (1.0 + np.exp(-x))

    for w, b in zip(self.weights, self.biases):
      x = sigmoid(np.dot(w, x.T).T + b)
    return x
Performs the forward pass.

    Args:
      x: 2-d array of size batch_size x image_size.

    Returns:
      A 2-d array of size batch_size x num_classes.
    
trax/tf_numpy/examples/mnist/model.py