Function stack

mlx/ops.cpp:1267–1287

Stack arrays along a new axis */

/** Stack arrays along a new axis */
array stack(
    const std::vector<array>& arrays,
    int axis,
    StreamOrDevice s /* = {} */) {
  if (arrays.empty()) {
    throw std::invalid_argument("[stack] No arrays provided for stacking");
  }
  if (!std::all_of(arrays.begin(), arrays.end(), [&](const auto& a) {
        return arrays[0].shape() == a.shape();
      })) {
    throw std::invalid_argument("[stack] All arrays must have the same shape");
  }
  auto normalized_axis =
      normalize_axis_index(axis, arrays[0].ndim() + 1, "[stack] ");
  std::vector<array> new_arrays;
  new_arrays.reserve(arrays.size());
  for (auto& a : arrays) {
    new_arrays.emplace_back(expand_dims(a, normalized_axis, s));
  }
  return concatenate(new_arrays, axis, s);
}

array stack(const std::vector<array>& arrays, StreamOrDevice s /* = {} */) {
  return stack(arrays, 0, s);