How to use slicing of data for NDArray structure?

[musram]

How to get something similar to output[:,0,:] if output is of type of NDArray? I am using for example int64_t X = logits_on_device->shape[0]; int64_t Y = logits_on_device->shape[1]; int64_t Z = logits_on_device->shape[2]; DLDataType dtype = logits_on_device->dtype;

      // Create a new NDArray for the slice with shape [X, Z]
      NDArray next_tokens_logits = NDArray::Empty({X, Z}, dtype, logits_on_device->device);

      // Copy data from 'logits_on_device' to 'slice'
      for (int64_t x = 0; x < X; ++x) {
        for (int64_t z = 0; z < Z; ++z) {
          // Calculate the index in the original and new array
          int original_index = (x * Y * Z) + z;  // Y is the size of the second dimension
          int new_index = x * Z + z;

          // Copy the element
          static_cast<float*>(next_tokens_logits->data)[new_index] =
              static_cast<float*>(logits_on_device->data)[original_index];
        }
      }

I have multiple slice so don't think its optimal to do this operation. How can these ops can be optimized? @

[ysh329]

Please open a new thread on https://discuss.tvm.ai/ for general support questions :)