Function mlx_expand_ellipsis

python/src/indexing.cpp:210–269

}

auto mlx_expand_ellipsis(const mx::Shape& shape, const nb::tuple& entries) {
  std::vector<nb::object> indices;

  // Go over all entries and note the position of ellipsis
  int non_none_indices_before = 0;
  int non_none_indices_after = 0;
  std::vector<nb::object> r_indices;
  int i = 0;
  bool has_ellipsis = false;

  // Start from dimension 0 till we hit an ellipsis
  for (; i < entries.size(); i++) {
    auto idx = entries[i];
    if (!is_valid_index_type(idx)) {
      throw std::invalid_argument(
          "Cannot index mlx array using the given type yet");
    }
    if (!nb::ellipsis().is(idx)) {
      indices.push_back(idx);
      non_none_indices_before += !idx.is_none();
    } else {
      has_ellipsis = true;
      break;
    }
  }

  // If we do hit an ellipsis, collect indices from the back
  for (int j = entries.size() - 1; j > i; j--) {
    auto idx = entries[j];
    if (!is_valid_index_type(idx)) {
      throw std::invalid_argument(
          "Cannot index mlx array using the given type yet");
    }
    if (nb::ellipsis().is(idx)) {
      throw std::invalid_argument(
          "An index can only have a single ellipsis (...)");
    }
    r_indices.push_back(idx);
    non_none_indices_after += !idx.is_none();
  }

  // Count up the number of non none indices
  int non_none_indices = non_none_indices_before + non_none_indices_after;

  // Expand ellipsis
  if (has_ellipsis) {
    for (int axis = non_none_indices_before;
         axis < shape.size() - non_none_indices_after;
         axis++) {
      indices.push_back(
          nb::slice(mx::ShapeElem{0}, shape[axis], mx::ShapeElem{1}));
      non_none_indices++;
    }
  }

  // Insert indices collected after the ellipsis
  indices.insert(indices.end(), r_indices.rbegin(), r_indices.rend());