func callAsFunction(_ x: MLXArray) -> MLXArray {
   var x = hiddenLayer(x)
   x = relu(x)
   x = outputLayer(x)
   return x
}