Repository: huggingface/candle
File: candle-core/src/tensor.rs (Rust)

Code before:
    /// - If inputs are 1D vectors (`[n]`), returns their scalar dot product.
    /// - If inputs are 2D matrices (`[m, n]` and `[m, n]`), returns the Frobenius inner product
    ///   (sum of element-wise products).
    /// - Panics if shapes are incompatible.
    ///
    /// # Example (vectors)
    /// ```rust
    /// let a = Tensor::new(&[1.0, 2.0, 3.0], &Device::Cpu)?;
    /// let b = Tensor::new(&[4.0, 5.0, 6.0], &Device::Cpu)?;
    /// let res = a.dot(&b)?;  // 1 * 4 + 2 * 5 + 3 * 6 = 32
    /// ```
    pub fn dot(&self, rhs: &Self) -> Result<Self> {
        if self.shape() != rhs.shape() {
            return Err(Error::ShapeMismatchBinaryOp {
                lhs: self.shape().clone(),
                rhs: rhs.shape().clone(),
                op: "dot",
            });
        }
        (self * rhs)
            .and_then(|ret| ret.sum_all())
    }

    /// Computes the **Frobenius norm** (L2 norm of all elements) of the tensor.
    /// - Output is `sqrt(sum(x^2))`.
    /// - Always returns a scalar (`[]` shape).
    ///
    /// # Example
    /// ```rust
    /// let t = Tensor::new(&[[3, 4], [0, 0]], &Device::Cpu)?;
    /// let norm = t.norm()?; // sqrt(3² + 4² + 0² + 0²) = 5 (exact integer)
    /// ```
    pub fn norm(&self) -> Result<Self> {
        self.sqr()
            .and_then(|x| x.sum_all())
            .and_then(|x| x.sqrt())
    }

    /// Performs strict matrix-vector multiplication (`[m, n] * [n] = [m]`).
    ///
    /// - If `self` is a matrix (`[m, n]`) and `rhs` is a vector (`[n]`), returns a vector (`[m]`).
    /// - **No broadcasting**: Panics if `self` is not 2D or if `rhs` is not 1D with matching size.
    ///
    /// # Example
    /// ```rust
    /// let mat = Tensor::randn(&[2, 3], &Device::Cpu)?;  // 2x3
    /// let vec = Tensor::randn(&[3], &Device::Cpu)?;     // 3
    /// let res = mat.mv(&vec)?;                          // Output shape [2]
    /// ```
    pub fn mv(&self, rhs: &Self) -> Result<Self> {
        // Strict shape checks

Review comment (bug):
Returning an int result on an int tensor seems pretty error prone as the output will be rounded + it's easy to do computations that would result in overflowing on u8 or even u32. I would suggest also being in line with the pytorch behavior here so bail on integer type tensors.

Code after:
    /// - If inputs are 1D vectors (`[n]`), returns their scalar dot product.
    /// - Panics if shapes are not compatible
    /// - Not supported for integer dtypes
    ///
    /// # Example (vectors)
    /// ```rust
    /// use candle_core::{Tensor, Device};
    /// let t1 = Tensor::new(&[1.0, 2.0, 3.0], &Device::Cpu)?;
    /// let t2 = Tensor::new(&[4.0, 5.0, 6.0], &Device::Cpu)?;
    /// let res = t1.dot(&t2)?;
    /// assert_eq!(res.to_scalar::<f64>()?, 32.);
    /// # Ok::<(), candle_core::Error>(())
    /// ```
    pub fn dot(&self, rhs: &Self) -> Result<Self> {
        if self.dims().len() != 1 || rhs.dims().len() != 1 {
            return Err(Error::ShapeMismatchBinaryOp {
                lhs: self.shape().clone(),
                rhs: rhs.shape().clone(),
                op: "dot",
            });
        }

        (self * rhs).and_then(|ret| ret.sum_all())
    }

    /// Computes the **Frobenius norm** (L2 norm of all elements) of the tensor.
    /// - Output is `sqrt(sum(x^2))`.
    /// - Always returns a scalar (`[]` shape).
    ///
    /// # Example
    /// ```rust
    /// use candle_core::{Tensor, Device};
    /// let t = Tensor::new(&[[3., 4.], [0., 0.]], &Device::Cpu)?;
    /// let norm = t.norm()?;
    /// assert_eq!(norm.to_scalar::<f64>()?, 5.);
    /// # Ok::<(), candle_core::Error>(())
    /// ```
    pub fn norm(&self) -> Result<Self> {
        if self.dtype().is_int() {
            bail!("norm not supported for integer dtypes");
        }

        self.sqr().and_then(|x| x.sum_all()).and_then(|x| x.sqrt())
    }

    /// Performs strict matrix-vector multiplication (`[m, n] * [n] = [m]`).
    ///
    /// - If `self` is a matrix (`[m, n]`) and `rhs` is a vector (`[n]`), returns a vector (`[m]`).
    /// - **No broadcasting**: Panics if `self` is not 2D or if `rhs` is not 1D with matching size.
    ///
    /// # Example