MLX
Loading...
Searching...
No Matches
indexing.h
Go to the documentation of this file.
1// Copyright © 2023-2024 Apple Inc.
2
3#pragma once
4
5#include <metal_stdlib>
6
7template <typename IdxT, int NIDX>
8struct Indices {
9 const array<const device IdxT*, NIDX> buffers;
10 const constant int* shapes;
11 const constant size_t* strides;
12 const int ndim;
13};
14
15template <typename IdxT>
16METAL_FUNC size_t offset_neg_idx(IdxT idx, size_t size) {
17 if (is_unsigned_v<IdxT>) {
18 return idx;
19 } else {
20 return (idx < 0) ? idx + size : idx;
21 }
22}
METAL_FUNC size_t offset_neg_idx(IdxT idx, size_t size)
Definition indexing.h:16
Definition indexing.h:8
const constant int * shapes
Definition indexing.h:10
const int ndim
Definition indexing.h:12
const constant size_t * strides
Definition indexing.h:11
const array< const device IdxT *, NIDX > buffers
Definition indexing.h:9