Function to_py_array

python/src/wrap.cpp:183–222

};

py::array to_py_array(const std::shared_ptr<const mlx::data::Array>& a) {
  const py::capsule free_when_done(new PyArrayPayload({a}), [](void* payload) {
    delete reinterpret_cast<PyArrayPayload*>(payload);
  });
  std::vector<int64_t> shape;
  std::vector<int64_t> stride;
  py::dtype pytype;
  switch (a->type()) {
    case mlx::data::ArrayType::Int8:
      pytype = py::dtype("b");
      break;
    case mlx::data::ArrayType::UInt8:
      pytype = py::dtype("B");
      break;
    case mlx::data::ArrayType::Int32:
      pytype = py::dtype("i4");
      break;
    case mlx::data::ArrayType::Int64:
      pytype = py::dtype("i8");
      break;
    case mlx::data::ArrayType::Float:
      pytype = py::dtype("f");
      break;
    case mlx::data::ArrayType::Double:
      pytype = py::dtype("d");
      break;
    default:
      throw std::runtime_error("internal error: unknown type");
  }
  auto ndim = a->shape().size();
  if (ndim > 0) {
    stride = std::vector<int64_t>(ndim);
    shape = a->shape();
    stride[ndim - 1] = a->itemsize();
    for (int i = ndim - 1; i > 0; i--) {
      stride[i - 1] = stride[i] * a->shape(i);
    }
  }
  return py::array(pytype, shape, stride, a->data(), free_when_done);
}

py::dict to_py_sample(const Sample& s) {
  py::dict res;