Rust Arrayfire - Using an array to index into another array

Question

For the last three days I've been trying to implement a loss function (cross entropy) in my machine learning project following the book "Neural Networks from Scratch by Harrison Kinsley".

So far I got an array with class predictions from a softmax activation layer:

[0.3298     0.3351     0.3351]
[0.3298     0.3351     0.3351]
[0.3298     0.3351     0.3351]
[0.3298     0.3351     0.3351]
...

And an array with target labels: [0, 2, 1, 1 ...]

To calculate loss I can take mean of all the correct confidences' negative natural logarithms. (-log(0.3298) + -log(0.3351) + -log(0.3351) + -log(0.3351)) / 4

But I am stuck on extracting these "correct confidence scores" from the class predictions array.
Is there any way to directly use the target labels array to do this using the arrayfire rust bindings, such as is possible with numpy:

predictions[range(l), target_labels]

I've tried using a vector of Seqs, but it seems inefficient if sample sizes go up.

Any help on this is greatly appreciated. (If this question is a duplicate then i am to dumb to find the original post)