MCPcopy Create free account
hub / github.com/ml-explore/mlx-data / Batch

Class Batch

mlx/data/buffer/Batch.h:11–35  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

9namespace buffer {
10
11class Batch : public Buffer {
12 public:
13 Batch(
14 const std::shared_ptr<Buffer>& op,
15 int64_t batch_size,
16 const std::unordered_map<std::string, double>& pad_values = {},
17 const std::unordered_map<std::string, int>& batch_dims = {});
18 Batch(
19 const std::shared_ptr<Buffer>& op,
20 const std::vector<int64_t>& batch_sizes,
21 const std::unordered_map<std::string, double>& pad_values = {},
22 const std::unordered_map<std::string, int>& batch_dims = {});
23
24 virtual Sample get(int64_t idx) const override;
25 virtual int64_t size() const override;
26
27 private:
28 std::shared_ptr<Buffer> op_;
29 int64_t batchSize_;
30 std::vector<int64_t> batchOffsets_;
31 std::vector<int64_t> batchSizes_;
32 std::unordered_map<std::string, double> padValues_;
33 std::unordered_map<std::string, int> batchDims_;
34 int64_t size_;
35};
36
37} // namespace buffer
38} // namespace data

Callers

nothing calls this directly

Calls

no outgoing calls

Tested by

no test coverage detected