Working with NDArray (cpp)

[erickim555]

I’m starting to work with DGL’s C++ code base, and one thing I’m very unclear about is how to work with the NDArray class, which is used all over the place.

For instance: how do I read and write elements in an NDArray? eg if I have an NDArray with shape=[2, 3], I’d like to:

(1) Read the entry at ndarray[0, 1]
(2) Write ndarray[0, 1] = 42.0f

Any tips on the above? Thanks!

[BarclayII]

In general you only need the following members:

• array->dtype: DLDataType: the data type of the array as a DLDataType structure from DLPack.
• array->ctx: DLContext: the device of the array as a DLContext structure from DLPack.
• array->shape: int64_t *: the shape.
• array->ndim: int: the number of dimensions
• data = array.Ptr<T>(): cast NDArray into a raw pointer of data type T.

So to answer your question, you will need to compute the offset of [0, 1] yourself and then read/write the memory:

offset = 0 * array->shape[1] + 1;
const int64_t* read = array.Ptr<int64_t>();    // read
int64_t value = read[offset];
int64_t* write = array.Ptr<int64_t>();         // write
write[offset] = value;

If you want to work with multiple types, your best option will be using ATEN_DTYPE_SWITCH macro together with templates.