Accessing raw array data from C++

[samrat]

Hello,

I would like to access the raw bytes of an array(using the C++ API). I've noticed that when I'm trying to do this with an array that was the result of a broadcast, the bytes returned don't reflect the values in the array:

mlx::core::array arr;

  // Evaluate the array to ensure data is available
  mlx::core::eval(arr);
  std::cout << "Array elements: ";
  for (auto it = reshaped.begin(); it != reshaped.end(); ++it) {
    std::cout << *it << " ";
  }

  std::cout << std::endl;

  std::cout << "contiguous: " << reshaped.flags().contiguous << std::endl;
  std::cout << "row contiguous: " << reshaped.flags().row_contiguous
            << std::endl;

  // Get raw pointer to data
  const void *data_ptr = reshaped.data<void>();

  if (data_ptr == nullptr) {
    return nx::nif::error(env, "Failed to get tensor data");
  }

  std::cout << "Data_ptr bytes: ";
  uint8_t *data_bytes = static_cast<uint8_t *>(const_cast<void *>(data_ptr));
  for (size_t i = 0; i < byte_size; i++) {
    std::cout << static_cast<int>(data_bytes[i]) << " ";
  }
  std::cout << std::endl;

Even when the flags.contiguous value is true(1), I get the following output:

Array elements: array(2, dtype=int32) array(2, dtype=int32) array(2, dtype=int32) array(2, dtype=int32)
      contiguous: 1
      row contiguous: 0
      Data_ptr bytes: 2 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0

As you can see, even though MLX can access 4 int32's with value 2, the bytes has just one 2 with the rest of the values 0.

Am I accessing the raw bytes of the array incorrectly? And is there a way to reliably get the bytes contiguously?

Thanks!

[samrat]

Looks like row_contiguous and col_contigious are the relevant flags(

mlx/mlx/io/load.cpp

Line 117 in cb431df

a = reshape(flatten(a), a.shape());

) -- in my case, both are 0.

I tried the reshape + flatten that the save code in load.cpp above does when the flags are false, but they don't change the flags for me 🤔

[awni]

Yes, I think this would be useful. For some context, we are working on an MLX backend for Nx.

That's cool!

I went with the ContiguousIterator approach for now.

It works like that:

    ContiguousIterator it(a);
    for (size_t elem = 0; elem < a.size(); elem +++) {
      a_ptr[it.loc];
      it.step();
    }

When printing out a tensor, we pass the bytes of the array & shape to the Elixir layer.

Why not use the MLX print routines instead?

Yes, I think this would be useful.

Can you say more about how you would use it / why you need it? It's still not clear to me what you are trying to do with the raw data.

[samrat]

Why not use the MLX print routines instead?
Can you say more about how you would use it / why you need it?

We're passing the raw data to create an Elixir representation of the tensor, once the array is eval'd(printing the tensor is just one use case-- sorry, should have been clearer :) ). Internally, we store a reference to an MLX array. But when returning a tensor to the Nx user, we'd return an Elixir data structure instead of the reference.

[polvalente]

@awni Hi! Nx maintainer here. This is especially useful when transferring data between different frameworks and when printing data using the Inspect protocol (think of it as the __repr__ method in Python).

Nx agrees on a contiguous row-major representation by default, and then we either obtain a copy of either a slice of the data (for Inspect) or all of the data that we then allocate as an ErlNifBinary (which is what represents a binary blob in Elixir from within the C interface).

Obtaining all of the data only happens when really needed. Either because someone requested to inspect the whole tensor, or because we're handing over the tensor to another backend (XLA, PyTorch...)

[awni]

Cool! We added a contiguous op in the latest MLX in the C++ back-end. So you can do:

auto a = some computation...;
a = contiguous(a);
eval(a);

auto ptr = a.data<void>(); // At this point should be row contiguous

[polvalente]

This is actually perfect for another functionality we have! Nx also has to_pointer and from_pointer, which can be used to transfer data ownership between Elixir and some other language. This is exactly what I was missing for that <3