Moonbit docs

///|
fn layer_norm_context(
  x : Tensor,
  weight : Tensor,
  bias : Tensor,
) -> AutogradContext? {
  match (result_context(x, weight), bias.context) {
    (Some(ctx), Some(bctx)) => {
      require_same_context(ctx, bctx)
      Some(ctx)
    }
    (Some(ctx), None) => Some(ctx)
    (None, Some(ctx)) => Some(ctx)
    (None, None) => None
  }
}

///|
pub fn Tensor::layer_norm(
  self : Tensor,
  weight : Tensor,
  bias : Tensor,
  eps : Double,
) -> Tensor {
  if eps < 0.0 {
    abort("layer_norm eps must not be negative")
  }
  if self.shape.length() == 0 {
    abort("layer_norm requires at least one dimension")
  }
  let hidden = self.shape[self.shape.length() - 1]
  if weight.shape != [hidden] || bias.shape != [hidden] {
    abort("layer_norm weight and bias must match the last dimension")
  }
  let rows = self.data.length() / hidden
  let out = Array::make(self.data.length(), 0.0)
  for row in 0.. {
        let node_ref = ctx.push_tape_node(
          self.shape,
          [tape_input(self), tape_input(weight), tape_input(bias)],
          LayerNormBackward::{
            input_data: self.data,
            weight_data: weight.data,
            eps,
          },
        )
        {
          data: out,
          shape: copy_ints(self.shape),
          requires_grad: true,
          context: Some(ctx),
          node_ref: Some(node_ref),
        }
      }
      None =>
        abort("differentiable tensor operation requires an autograd context")
    }
  } else {
    Tensor::from_array(out, self.shape)
  }
}