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

Function batch

mlx/data/Array.cpp:465–498  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

463}
464
465std::shared_ptr<Array> batch(
466 const std::vector<std::shared_ptr<Array>>& arrs,
467 double pad_value) {
468 core::BatchShape batch_shape;
469 auto ndim = arrs.front()->ndim();
470 auto type = arrs.front()->type();
471 for (int i = 0; i < arrs.size(); i++) {
472 if (arrs[i]->type() != type) {
473 throw std::runtime_error(
474 "Array: unexpected different types of arrays in batch");
475 }
476 batch_shape.add(arrs[i]->shape());
477 }
478 auto stride = std::vector<int64_t>(ndim);
479 int64_t item_stride = 1;
480 for (int dim = ndim - 1; dim >= 0; dim--) {
481 stride[dim] = item_stride;
482 item_stride *= batch_shape[dim + 1];
483 }
484 auto res = std::make_shared<Array>(type, batch_shape.shape());
485 res->fill(pad_value);
486 for (int i = 0; i < arrs.size(); i++) {
487 auto arr = arrs[i];
488 ARRAY_DISPATCH(
489 arr,
490 array_copy_linear_to_strided,
491 res->data(),
492 i * item_stride,
493 arr->data(),
494 arr->shape(),
495 stride);
496 }
497 return res;
498}
499
500std::shared_ptr<Array> batch(
501 const std::vector<std::shared_ptr<Array>>& arrs,

Callers 1

merge_batchFunction · 0.50

Calls 8

ndimMethod · 0.80
typeMethod · 0.80
fillMethod · 0.80
checkdimMethod · 0.80
sizeMethod · 0.45
addMethod · 0.45
shapeMethod · 0.45
dataMethod · 0.45

Tested by

no test coverage detected