Implementation:LaurentMazare Tch rs Tensor Ops
| Knowledge Sources | |
|---|---|
| Domains | Tensor_Operations, Operator_Overloading |
| Last Updated | 2026-02-08 00:00 GMT |
Overview
Concrete tool for using standard Rust arithmetic operators on tensors provided by the tch-rs library.
Description
The Tensor_Ops module implements Rust's standard operator traits (Add, Sub, Mul, Div, Neg, and their Assign variants) for the Tensor type. It uses four macros to generate all necessary combinations of owned and borrowed tensor operands, as well as scalar operands of types i32, i64, f32, and f64. The module also implements PartialEq for element-wise tensor comparison. This allows natural arithmetic syntax such as a + b, tensor * 2.0, and tensor -= other in Rust code.
Usage
These operator implementations are available automatically when the Tensor type is in scope. No additional trait imports are required for the standard arithmetic operators.
Code Reference
Source Location
- Repository: LaurentMazare_Tch_rs
- File: src/tensor/ops.rs
- Lines: 1-307
Signature
The module defines four code-generation macros and then invokes them for each operator:
// Tensor-Tensor operations (generates Tensor op Tensor, Tensor op &Tensor,
// &Tensor op &Tensor, &Tensor op Tensor)
macro_rules! impl_op {
($trait:ident, $func:ident, $op:ident) => { ... }
}
// Scalar-Tensor operations (generates i32/i64/f32/f64 op Tensor and op &Tensor)
macro_rules! impl_op_basic {
($trait:ident, $func:ident, $op:ident, $rev:ident) => { ... }
}
// In-place Tensor-Tensor operations (generates Tensor op= Tensor and Tensor op= &Tensor)
macro_rules! impl_op_assign {
($trait:ident, $func:ident, $op:ident) => { ... }
}
// In-place scalar operations (generates Tensor op= i32/i64/f32/f64)
macro_rules! impl_op_assign_basic {
($trait:ident, $func:ident, $op:ident) => { ... }
}
Additionally, Tensor + Scalar and &Tensor + Scalar are implemented directly using a generic S: Into<Scalar> bound for Add, Sub, Mul, and Div.
Negation and equality are implemented explicitly:
impl Neg for Tensor {
type Output = Tensor;
fn neg(self) -> Tensor { self.f_neg().unwrap() }
}
impl Neg for &Tensor {
type Output = Tensor;
fn neg(self) -> Tensor { self.f_neg().unwrap() }
}
impl PartialEq for Tensor {
fn eq(&self, other: &Tensor) -> bool { ... }
}
Import
use tch::Tensor; // Operator traits are auto-available
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| self | Tensor or &Tensor | Yes | Left-hand operand |
| rhs | Tensor, &Tensor, i32, i64, f32, f64, or S: Into<Scalar> | Yes | Right-hand operand |
Outputs
| Name | Type | Description |
|---|---|---|
| result | Tensor | Result of the arithmetic operation (for binary ops) |
| result | bool | Result of equality comparison (for PartialEq) |
Generated Operator Combinations
The macro invocations produce the following operator implementations:
| Operator | Macro | Underlying Method | Reversal Function |
|---|---|---|---|
| Add (Tensor + Tensor) | impl_op! | g_add | n/a |
| Add (scalar + Tensor) | impl_op_basic! | g_add_scalar | id (identity) |
| AddAssign (Tensor += Tensor) | impl_op_assign! | g_add_ | n/a |
| AddAssign (Tensor += scalar) | impl_op_assign_basic! | g_add_scalar_ | n/a |
| Mul (Tensor * Tensor) | impl_op! | g_mul | n/a |
| Mul (scalar * Tensor) | impl_op_basic! | g_mul_scalar | id (identity) |
| MulAssign (Tensor *= Tensor) | impl_op_assign! | g_mul_ | n/a |
| MulAssign (Tensor *= scalar) | impl_op_assign_basic! | g_mul_scalar_ | n/a |
| Div (Tensor / Tensor) | impl_op! | g_div | n/a |
| Div (scalar / Tensor) | impl_op_basic! | g_div_scalar | inv (reciprocal) |
| DivAssign (Tensor /= Tensor) | impl_op_assign! | g_div_ | n/a |
| DivAssign (Tensor /= scalar) | impl_op_assign_basic! | g_div_scalar_ | n/a |
| Sub (Tensor - Tensor) | impl_op! | g_sub | n/a |
| Sub (scalar - Tensor) | impl_op_basic! | g_sub_scalar | neg (negate) |
| SubAssign (Tensor -= Tensor) | impl_op_assign! | g_sub_ | n/a |
| SubAssign (Tensor -= scalar) | impl_op_assign_basic! | g_sub_scalar_ | n/a |
The rev parameter in impl_op_basic! handles the non-commutative case: when computing scalar - tensor, the macro computes tensor.g_sub_scalar(scalar) and then applies neg to reverse the operand order. Similarly, scalar / tensor computes tensor.g_div_scalar(scalar) and applies inv (reciprocal via pow_tensor_scalar(-1)).
Usage Examples
use tch::Tensor;
let a = Tensor::from_slice(&[1.0f32, 2.0, 3.0]);
let b = Tensor::from_slice(&[4.0f32, 5.0, 6.0]);
// Tensor-Tensor arithmetic
let c = &a + &b; // [5.0, 7.0, 9.0]
let d = &a * &b; // [4.0, 10.0, 18.0]
let e = &b - &a; // [3.0, 3.0, 3.0]
let f = &b / &a; // [4.0, 2.5, 2.0]
// Tensor-Scalar arithmetic
let g = &a + 10.0; // [11.0, 12.0, 13.0]
let h = &a * 2.0; // [2.0, 4.0, 6.0]
let i = 1.0 - &a; // [0.0, -1.0, -2.0]
// In-place operations
let mut t = Tensor::from_slice(&[1.0f32, 2.0, 3.0]);
t += &a; // [2.0, 4.0, 6.0]
t *= 2.0; // [4.0, 8.0, 12.0]
// Negation
let neg_a = -&a; // [-1.0, -2.0, -3.0]
// Equality
let eq = a == b; // false