Method eval_cpu

mlx/backend/cpu/softmax.cpp:117–168

} // namespace

void Softmax::eval_cpu(const std::vector<array>& inputs, array& out) {
  assert(inputs.size() == 1);

  // Make sure that the last dimension is contiguous
  auto set_output = [s = stream(), &out](const array& x) {
    if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) {
      if (x.is_donatable()) {
        out.copy_shared_buffer(x);
      } else {
        out.set_data(
            allocator::malloc(x.data_size() * x.itemsize()),
            x.data_size(),
            x.strides(),
            x.flags());
      }
      return x;
    } else {
      array x_copy = contiguous_copy_cpu(x, s);
      out.copy_shared_buffer(x_copy);
      return x_copy;
    }
  };

  auto in = set_output(inputs[0]);

  switch (in.dtype()) {
    case float32:
      softmax<float, float>(in, out, stream());
      break;
    case float16:
      if (precise_) {
        softmax<float16_t, float>(in, out, stream());
      } else {
        softmax<float16_t, float16_t>(in, out, stream());
      }
      break;
    case bfloat16:
      if (precise_) {
        softmax<bfloat16_t, float>(in, out, stream());
      } else {
        softmax<bfloat16_t, bfloat16_t>(in, out, stream());
      }
      break;
    case float64:
      softmax<double, double>(in, out, stream());
      break;
    default:
      throw std::runtime_error(
          "[softmax] Only defined for floating point types.");
      break;
  }
}

} // namespace mlx::core