Function tch_utils::ops_2d::dot_product_2d

pub fn dot_product_2d(a: &Tensor, b: &Tensor) -> Tensor

Compute the vector dot product between two tensors for each pixel.

Arguments

a: Tensor - The first tensor [N, C, H, W] b: Tensor - The second tensor [N, C, H, W]

Returns

Tensor - The dot product of the two tensors [N, 1 ,H, W]

Example

let a = Tensor::of_slice(&[1.0, 2.0, 3.0, 4.0]).view([1, 2, 2, 1]);
let b = Tensor::of_slice(&[1.0, 2.0, 3.0, 4.0]).view([1, 2, 2, 1]);
let dot = dot_product_2d(&a, &b);
assert!(dot.equal(&Tensor::of_slice(&[10.0, 20.0]).view([1, 1, 2, 1])));