ModularM
Modular•2y ago•
1 reply
pascalpolygon

Efficient Mojo Tensor To Numpy Array Conversion

Hello,

I am attempting to convert a Mojo Tensor to a Numpy array without looping or at least with sub-polynomial time.

My approach is to get the pointer to the Tensor and use Python's ctypes.from_address()

But it doesn't work - the output np_array is not the same as the tensor...maybe Tensor has some header bytes? Or the pointer does not point to a contiguous memory block?

I would like to understand the issue with this code and get some inputs on how to achieve fast conversions for potentially large Mojo Tensors.

Here is the code:

def main():
    var ctypes = Python.import_module("ctypes")
    var np = Python.import_module("numpy")

    var mojo_tensor =  Tensor[DType.int32](
        TensorShape(4, 4),
        1, 0, 0, 0,
        0, 1, 0, 0,
        0, 0, 1, 0,
        0, 0, 0, 1,
    )

    var mojo_tensor_ptr = mojo_tensor.data().__int__()
  
    var num_elements = mojo_tensor.num_elements().__int__()
    var tensor_type = mojo_tensor.type

    
    #Load the hopefully contiguous tensor into a ctypes array
    var int32_array = (ctypes.c_int32 * num_elements).from_address(mojo_tensor_ptr.__int__())
    var np_array = np.ctypeslib.as_array(int32_array).astype(np.int32)

    print("Numpy array: ", np_array)


Thank you.
Was this page helpful?