burn-jit 0.16.1

Generic backend that can be compiled just-in-time to any shader language target
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

use crate::{element::JitElement, ops::numeric::empty_device, tensor::JitTensor, JitRuntime};
use burn_tensor::Shape;

/// Creates an empty output tensor with matmul output shape
pub fn init_matmul_output<R: JitRuntime, E: JitElement>(
    lhs: &JitTensor<R>,
    rhs: &JitTensor<R>,
) -> JitTensor<R> {
    empty_device::<R, E>(lhs.client.clone(), lhs.device.clone(), shape_out(lhs, rhs))
}

pub(crate) fn shape_out<R: JitRuntime>(lhs: &JitTensor<R>, rhs: &JitTensor<R>) -> Shape {
    let ndims = lhs.shape.num_dims();
    let mut shape_out = vec![0; ndims];
    lhs.shape
        .dims
        .iter()
        .zip(rhs.shape.dims.iter())
        .enumerate()
        .for_each(|(index, (dim_lhs, dim_rhs))| {
            shape_out[index] = usize::max(*dim_lhs, *dim_rhs);
        });
    shape_out[ndims - 2] = lhs.shape.dims[ndims - 2];
    shape_out[ndims - 1] = rhs.shape.dims[ndims - 1];
    Shape::from(shape_out)
}