burn_tensor/tensor/linalg/
outer.rs

1use crate::backend::Backend;
2use crate::tensor::{BasicOps, Tensor};
3use crate::{AsIndex, Numeric};
4
5/// Computes the outer product for the last columns of 2 tensors.
6///
7/// See also: [`outer_dim`].
8///
9/// # Arguments
10/// - `lhs`: the "row" tensor, with shape ``[..., i]``.
11/// - `rhs`: the "col" tensor, with shape ``[..., j]``.
12///
13/// # Returns
14///
15/// A tensor of rank `R = D + 1`, where:
16///
17/// ``
18/// result[..., i, j] = lhs[..., i] * rhs[..., j]
19/// ``
20pub fn outer<B: Backend, const D: usize, const R: usize, K>(
21    lhs: Tensor<B, D, K>,
22    rhs: Tensor<B, D, K>,
23) -> Tensor<B, R, K>
24where
25    K: BasicOps<B> + Numeric<B>,
26{
27    outer_dim(lhs, rhs, -1)
28}
29
30/// Computes the outer product along a specific dimension, broadcasting over others.
31///
32/// For the given `dim`, computes the outer product of elements along that dimension,
33/// expanding it into two dimensions of size ``M × N`` at positions ``(dim, dim + 1)``.
34///
35/// # Arguments
36///
37/// - `lhs`: left operand, the "row" tensor, with size `M` at dimension `dim`.
38/// - `rhs`: right operand, the "col" tensor, with size `N` at dimension `dim`.
39/// - `dim`: dimension to compute the outer product along (supports negative indexing).
40///
41/// # Returns
42///
43/// A tensor of rank `R = D + 1`, where:
44///
45/// ``
46/// result[..., i, j, ...] = lhs[..., i, ...] * rhs[..., j, ...]
47/// ``
48//
49// Notes:
50// - For large batched inputs, `x_col.matmul(y_row)` *might* be more performant
51//   than broadcasted elemwise multiply; benchmarking needed to confirm.
52pub fn outer_dim<B: Backend, const D: usize, const R: usize, Dim: AsIndex, K>(
53    lhs: Tensor<B, D, K>,
54    rhs: Tensor<B, D, K>,
55    dim: Dim,
56) -> Tensor<B, R, K>
57where
58    K: BasicOps<B> + Numeric<B>,
59{
60    assert_eq!(
61        R,
62        D + 1,
63        "`outer` with D={D} expects R={} (got R={R})",
64        D + 1
65    );
66    let dim = dim.expect_dim_index(D);
67
68    // (..., i, 1, ...)
69    let x = lhs.unsqueeze_dim::<R>(dim + 1);
70
71    // (..., 1, j, ...)
72    let y = rhs.unsqueeze_dim::<R>(dim);
73
74    // (..., i, j, ...)
75    x * y
76}
burn_tensor/tensor/linalg/outer.rs

burn_tensor/tensor/linalg/
outer.rs
