Function categorical 

pub fn categorical<'a>(
    logits: &Array,
    axis: impl Axis,
    key: impl Into<Option<&'a Array>>,
) -> Array

Indices (uint32) drawn from the categorical distributions whose unnormalised log probabilities run along one axis of logits (..: the flattened array), as MLX: the argmax of logits plus Gumbel noise. The result has logits’ shape without axis.

use nuvec_core::{array, random};

// One certain outcome per row.
let logits = array![[0.0, -1e9, -1e9], [-1e9, -1e9, 0.0]];
let draws = random::categorical(&logits, -1, Some(&random::key(7)));
assert_eq!(draws.to_vec::<u32>().unwrap(), [0, 2]);