// nn.mbt

///|
/// Linear (fully connected) neural network layer.
pub struct Linear {
  weight : Tensor
  bias : Tensor?
}

///|
/// Create a new Linear layer with random weights (Xavier initialized) and zero bias.
pub fn Linear::new(
  in_features : Int,
  out_features : Int,
  requires_grad? : Bool = true,
) -> Linear {
  if in_features <= 0 || out_features <= 0 {
    panic()
  }
  let limit = (6.0 / (in_features + out_features).to_double()).sqrt()
  let weight_data = Array::make(in_features * out_features, 0.0)

  let rng = Random::new(42 + in_features * 7 + out_features * 13)
  let size = in_features * out_features
  for i in 0.. Tensor {
  let out = x.matmul(self.weight)
  match self.bias {
    Some(b) => out + b
    None => out
  }
}

///|
/// Get parameters (weight, bias) of the Linear layer.
pub fn Linear::parameters(self : Linear) -> Array[Tensor] {
  match self.bias {
    Some(b) => [self.weight, b]
    None => [self.weight]
  }
}

///|
/// Mean Squared Error (MSE) loss between prediction and target.
pub fn mse_loss(pred : Tensor, target : Tensor) -> Tensor {
  if pred.shape != target.shape {
    panic()
  }
  let size = pred.data.length()
  if size == 0 {
    panic()
  }
  let mut sum_val = 0.0
  for i in 0..