Function init_mlx_data_buffer

python/src/wrap_buffer.cpp:38–432

using namespace mlx::data;

void init_mlx_data_buffer(py::module& m) {
  py::class_<PyBufferIterator>(m, "BufferIterator")
      .def(
          "__iter__",
          [](PyBufferIterator& it) -> PyBufferIterator& { return it; })
      .def("__next__", &PyBufferIterator::next);

  auto buffer_class =
      py::class_<
          Buffer,
          std::unique_ptr<Buffer, mlx::pybind::NogilDeleter<Buffer>>>(
          m, "Buffer")
          .def("__iter__", [](Buffer& buf) { return PyBufferIterator(buf); })
          .def("size", &Buffer::size, py::call_guard<py::gil_scoped_release>())
          .def(
              "__len__",
              &Buffer::size,
              py::call_guard<py::gil_scoped_release>())
          .def(
              "__getitem__",
              [](const Buffer& b, int64_t idx) {
                Sample sample;
                {
                  py::gil_scoped_release release;
                  idx = (idx < 0) ? idx + b.size() : idx;
                  sample = b.get(idx);
                }
                py::dict pysample;
                for (auto& item : sample) {
                  pysample[py::str(item.first)] =
                      mlx::pybind::to_py_array(item.second);
                }
                return pysample;
              },
              py::arg("idx"))
          .def(
              "__repr__",
              [](const Buffer& b) {
                std::stringstream msg;
                msg << "Buffer(size=" << b.size() << ", keys={";
                if (b.size() > 0) {
                  Sample first = b.get(0);
                  int cnt = 0;
                  for (auto& item : first) {
                    if (cnt++ > 0) {
                      msg << ", ";
                    }
                    msg << "'" << item.first << "'";
                  }
                }
                msg << "})";

                return msg.str();
              })
          .def(
              "batch",
              [](const Buffer& b,
                 const std::variant<int64_t, std::vector<int64_t>>& batch_size,