// softmax.mbt

///|
fn validate_softmax_dim(shape : Array[Int], dim : Int) -> Unit {
  if dim < 0 || dim >= shape.length() || shape[dim] == 0 {
    panic()
  }
}

///|
fn softmax_data(tensor : Tensor, dim : Int) -> Array[Double] {
  validate_softmax_dim(tensor.shape, dim)
  let data = Array::make(tensor.data.length(), 0.0)
  let width = tensor.shape[dim]
  let stride = tensor.strides[dim]
  for i in 0.. max_value {
        max_value = value
      }
    }
    let mut denominator = 0.0
    for j in 0.. Array[Double] {
  validate_softmax_dim(tensor.shape, dim)
  let data = Array::make(tensor.data.length(), 0.0)
  let width = tensor.shape[dim]
  let stride = tensor.strides[dim]
  for i in 0.. max_value {
        max_value = value
      }
    }
    let mut denominator = 0.0
    for j in 0.. Tensor {
  let data = softmax_data(self, dim)
  let out = Tensor::new(data, self.shape, requires_grad=self.requires_grad)
  if self.requires_grad {
    out.creator = Some(Softmax(self, dim))
  }
  out
}

///|
/// Apply logarithmic softmax along a dimension.
pub fn Tensor::log_softmax(self : Tensor, dim : Int) -> Tensor {
  let data = log_softmax_data(self, dim)
  let out = Tensor::new(data, self.shape, requires_grad=self.requires_grad)
  if self.requires_grad {
    out.creator = Some(LogSoftmax(self, dim))
  }
  out
}

///|
/// Return the log-sum-exp of a tensor as a scalar.
pub fn Tensor::logsumexp(self : Tensor) -> Tensor {
  if self.data.length() == 0 {
    panic()
  }
  let mut max_value = self.data[0]
  for value in self.data {
    if value > max_value {
      max_value = value
    }
  }
  let mut total = 0.0
  for value in self.data {
    total = total + exp(value - max_value)
  }
  Tensor::new([ln(total) + max_value], [], requires_grad=false)
}