1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

use anyhow::Error;
use rand::{Rng, SeedableRng};

use crate::{builtins::Arguments, Tensor};

pub fn random(args: &Arguments) -> Result<Tensor, Error> {
    let count: usize = args.parse_or_default("amount", 1)?;

    let rng = rand::thread_rng();
    random_tensor(count, rng)
}

pub fn seeded_random(args: &Arguments, seed: u64) -> Result<Tensor, Error> {
    let count: usize = args.parse_or_default("amount", 1)?;

    let rng = rand::rngs::SmallRng::seed_from_u64(seed);
    random_tensor(count, rng)
}

fn random_tensor(count: usize, mut rng: impl Rng) -> Result<Tensor, Error> {
    let mut numbers: Vec<u32> = Vec::new();
    for _ in 0..count {
        numbers.push(rng.gen());
    }
    Ok(Tensor::new(&numbers, &[1, count]))
}