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

Function pad

mlx/data/Array.cpp:254–299  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

252}
253
254std::shared_ptr<Array> pad(
255 const std::shared_ptr<const Array>& arr,
256 int dim,
257 int64_t lpad,
258 int64_t rpad,
259 double value) {
260 if (lpad == 0 && rpad == 0) {
261 return std::make_shared<Array>(arr);
262 }
263 if (lpad < 0 || rpad < 0) {
264 throw std::runtime_error("Array: pad must be positive");
265 }
266 auto ndim = arr->ndim();
267 if (dim < 0 || dim >= ndim) {
268 throw std::runtime_error("Array: dim out of range");
269 }
270 auto shape = arr->shape();
271 int64_t lchunksize = lpad;
272 int64_t chunksize = shape[dim];
273 int64_t rchunksize = rpad;
274 int64_t n = 1;
275 for (int d = dim + 1; d < ndim; d++) {
276 lchunksize *= shape[d];
277 chunksize *= shape[d];
278 rchunksize *= shape[d];
279 }
280 for (int d = 0; d < dim; d++) {
281 n *= shape[d];
282 }
283 shape[dim] += (lpad + rpad);
284 auto res = std::make_shared<Array>(arr->type(), shape);
285 res->fill(value);
286
287 ARRAY_DISPATCH(
288 arr,
289 array_pad,
290 res->data(),
291 arr->data(),
292 lchunksize,
293 chunksize,
294 rchunksize,
295 n,
296 arr->itemsize());
297
298 return res;
299}
300
301std::shared_ptr<Array> slice(
302 const std::shared_ptr<const Array>& arr,

Callers 1

apply_keyMethod · 0.50

Calls 6

ndimMethod · 0.80
typeMethod · 0.80
fillMethod · 0.80
itemsizeMethod · 0.80
shapeMethod · 0.45
dataMethod · 0.45

Tested by

no test coverage detected