Function mlx_random_categorical_shape 

pub unsafe extern "C" fn mlx_random_categorical_shape(
    res: *mut mlx_array,
    logits: mlx_array,
    axis: c_int,
    shape: *const c_int,
    shape_num: usize,
    key: mlx_array,
    s: mlx_stream,
) -> c_int