How to reinterpret_cast a matrix?
Closed this issue · 2 comments
Dear @dfm , your tutorial is excellent. But I am not familar with c++. I have a naive question.
You said that how to receive input values by:
#include <cstdint> // int64_t
template <typename T>
void cpu_kepler(void *out, const void **in) {
const std::int64_t size = *reinterpret_cast<const std::int64_t *>(in[0]);
const T *mean_anom = reinterpret_cast<const T *>(in[1]);
const T *ecc = reinterpret_cast<const T *>(in[2]);
}
However, if one of my input values is a matrix, how can I reinterpret_cast
it?
Thanks.
The parameters will always be raw pointers, even if they are defined as matrices or tensors in Python. So if you need more structure, you'll have to pass all the necessary parameters (e.g. strides) to your function as extra input parameters, and then use those to index your pointer appropriately. The example here works with any shape of input parameters, but just treats them as flat and aligned.
The logic for how this is handled is all in the XLA translation rule:
extending-jax/src/kepler_jax/kepler_jax.py
Lines 59 to 130 in 77c3466
Hope this helps!
Dear @dfm , thank you very much.