from flax import linen as nn

class FeedForward(nn.Module):
  hidden_dim: int

  def setup(self):
    # Define a dense layer with hidden_dim units
    self.dense = nn.Dense(self.hidden_dim) 
    # Define an output layer with 10 units 
    self.output = nn.Dense(10)  

  def __call__(self, inputs):
    # Apply ReLU activation to the dense layer
    x = nn.relu(self.dense(inputs))  
    # Return the output
    return self.output(x)

        __
        __