Add DeviceType to Array2D type annotation

[horizon-blue]

Recreating #21 because I accidentally closes it and GitHub doesn't let me reopen the PR..

---

(Closes MET-14)

Summary of Changes

Similar to raw pointers, our current Array2D type does not record where the underlying data comes from, as the memory accessing pattern are roughly the same on CPU & GPU. However, mentally tacking the location of the memory can be error-prone. In addition, many other Python array/tensor libraries store the device type explicitly, and we won't be able easily convert to them without knowing where our memory is.

As such, I'm adding a DeviceType template parameter to our Array2D to keep track of the location of the memory. With this change, we are finally able to take CPU/GPU buffers from Python side and return them correctly without running into segfault.

Test Plans

You can find examples of creating Array2D from CPU/GPU memory buffers with numpy and JAX ih the included test_utils.py.

As always, to run all the tests:

pixi run test

[linear]

MET-14 Copy-free data movement to/from Python

and ideally, we should have something that works for both device & host.

One idea to achieve this is to use DLPack which has been widely adopted by ML & scientific computing libraries.

[horizon-blue]

Update: rename DeviceType::CPU/DeviceType::GPU to MemoryLocation::HOST/MemoryLocation::DEVICE based on @mugamma's suggestion.

[horizon-blue]

I'm going to merge this PR now because the original one has been approved