ndrs 0.5.0

A tensor library with GPU support
Documentation 
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
27
28

// ndrs/src/backend/cpu/fill.rs
use super::elementwise::*;
use crate::device::Device;
use crate::dtype::DTypeMapping;
use crate::tensor::Tensor;
use crate::view::TensorViewOps;
use crate::view::{ArcTensorView, AsView, RcTensorView};
use anyhow::Result;

impl RcTensorView {
    /// 填充 CPU 视图为常量值
    pub fn fill_cpu<T: bytemuck::Pod + DTypeMapping>(&mut self, value: T) -> Result<()> {
        let scalar = Tensor::from_vec(vec![value], vec![1], Device::Cpu)?;
        let scalar_view = RcTensorView::new(scalar.into_rc());
        let broadcast = scalar_view.broadcast_to(self.shape())?;
        RcTensorView::elementwise::<T, _>(self, &[&broadcast], |_| value)
    }
}

impl ArcTensorView {
    /// 填充 CPU 视图为常量值
    pub fn fill_cpu<T: bytemuck::Pod + DTypeMapping>(&mut self, value: T) -> Result<()> {
        let scalar = Tensor::from_vec(vec![value], vec![1], Device::Cpu)?;
        let scalar_view = ArcTensorView::new(scalar.into_arc());
        let broadcast = scalar_view.broadcast_to(self.shape())?;
        ArcTensorView::elementwise::<T, _>(self, &[&broadcast], |_| value)
    }
}