Trait TensorRngPrims

Source

pub trait TensorRngPrims<Alg: Algebra> {
    type Plan;
    type Context;

    // Required methods
    fn plan(
        ctx: &mut Self::Context,
        desc: &RngPrimsDescriptor,
        shapes: &[&[usize]],
    ) -> Result<Self::Plan>;
    fn execute(
        ctx: &mut Self::Context,
        plan: &Self::Plan,
        generator: &mut Generator,
        output: &mut Tensor<Alg::Scalar>,
    ) -> Result<()>;
    fn has_rng_support(desc: RngPrimsDescriptor) -> bool;
}

Expand description

Tensor RNG execution family.

The family is plan/execute based like the rest of tenferro-prims, but the plan is intentionally small: it only records the descriptor and the output shape so the execution side can validate the destination tensor before writing into it.

§Examples

use tenferro_algebra::Standard;
use tenferro_device::Generator;
use tenferro_prims::{CpuBackend, CpuContext, RngPrimsDescriptor, TensorRngPrims};
use tenferro_tensor::{MemoryOrder, Tensor};

let mut ctx = CpuContext::new(1);
let mut generator = Generator::cpu(1234);
let desc = RngPrimsDescriptor::Uniform;
let mut output = Tensor::<f64>::zeros(
    &[8],
    tenferro_device::LogicalMemorySpace::MainMemory,
    MemoryOrder::ColumnMajor,
).unwrap();
let plan = <CpuBackend as TensorRngPrims<Standard<f64>>>::plan(
    &mut ctx,
    &desc,
    &[output.dims()],
).unwrap();
<CpuBackend as TensorRngPrims<Standard<f64>>>::execute(
    &mut ctx,
    &plan,
    &mut generator,
    &mut output,
).unwrap();