Implementation:LaurentMazare Tch rs Tensor Mod

Knowledge Sources	LaurentMazare_Tch_rs
Domains	Tensor Operations, Shape Abstraction, Module Organization
Last Updated	2026-02-08 00:00 GMT

Overview

The tensor/mod.rs file is the root of the tensor module, defining the Shape trait for flexible tensor shape specification, padding utilities, convenience methods for common tensor operations, and From trait implementations for creating tensors from Rust primitives and slices.

Description

This file serves as the organizational root for the tensor submodule, declaring and re-exporting from child modules: convert, display, index, iter, npy, ops, and safetensors. It re-exports key types from wrappers::tensor (Tensor, Reduction, NoGradGuard, etc.) and from index (IndexOp, NewAxis, TensorIndexer).

The Shape trait is a central abstraction that converts various Rust types to Box<[i64]> for use with tensor operations like view. Implementations are provided via the impl_shape! macro for [i64; N] arrays (N = 0..6), plus manual implementations for () (scalar/empty shape), &[i64] (slices), i64, usize, i32, and tuples of 1 to 4 i64 values.

The file adds several convenience methods to Tensor:

• f_view / view: Reshape using any Shape-implementing type
• f_zero_pad1d / zero_pad1d: Constant zero padding for 3D tensors (left, right)
• f_zero_pad2d / zero_pad2d: Constant zero padding for 4D tensors (left, right, top, bottom)
• to_kind / f_to_kind: Cast to a specified Kind
• nll_loss: Negative log-likelihood loss
• cross_entropy_for_logits: Cross-entropy loss from logits
• accuracy_for_logits: Classification accuracy computation
• random_batch / random_batch2: Random batch sampling
• to_device / f_to_device: Move tensor to a device
• avg_pool2d_default / max_pool2d_default: Default pooling operations
• flat_view: Flatten all dimensions except batch
• onehot: One-hot encoding
• copy: Deep copy
• from_slice2: Create 2D tensor from nested slices
• to_mkldnn: Convert to MKLDNN format

From trait implementations allow implicit conversion from &[T] (slices) and scalar T values to Tensor.

Usage

Use the Shape trait to pass diverse Rust types to tensor reshape operations. Use the convenience methods for common deep learning operations like loss computation, batch sampling, and padding.

Code Reference

Source Location

• Repository: LaurentMazare_Tch_rs
• File: src/tensor/mod.rs

Signature

// Shape trait
pub trait Shape {
    fn to_shape(&self) -> Box<[i64]>;
}
// Implemented for: [i64; 0..6], (), &[i64], i64, usize, i32,
//                  (i64,), (i64, i64), (i64, i64, i64), (i64, i64, i64, i64)

// View methods
impl Tensor {
    pub fn f_view<T: Shape>(&self, s: T) -> Result<Tensor, TchError>;
    pub fn view<T: Shape>(&self, s: T) -> Tensor;
}

// Padding
impl Tensor {
    pub fn f_zero_pad1d(&self, left: i64, right: i64) -> Result<Tensor, TchError>;
    pub fn zero_pad1d(&self, left: i64, right: i64) -> Tensor;
    pub fn f_zero_pad2d(&self, left: i64, right: i64, top: i64, bottom: i64) -> Result<Tensor, TchError>;
    pub fn zero_pad2d(&self, left: i64, right: i64, top: i64, bottom: i64) -> Tensor;
}

// Type and device conversion
impl Tensor {
    pub fn to_kind(&self, kind: Kind) -> Tensor;
    pub fn f_to_kind(&self, kind: Kind) -> Result<Tensor, TchError>;
    pub fn to_device(&self, device: Device) -> Tensor;
    pub fn f_to_device(&self, device: Device) -> Result<Tensor, TchError>;
}

// Loss and accuracy
impl Tensor {
    pub fn nll_loss(&self, targets: &Tensor) -> Tensor;
    pub fn cross_entropy_for_logits(&self, targets: &Tensor) -> Tensor;
    pub fn accuracy_for_logits(&self, targets: &Tensor) -> Tensor;
}

// Utilities
impl Tensor {
    pub fn random_batch(&self, batch_size: i64) -> Tensor;
    pub fn random_batch2(t1: &Tensor, t2: &Tensor, batch_size: i64) -> (Tensor, Tensor);
    pub fn flat_view(&self) -> Tensor;
    pub fn onehot(&self, labels: i64) -> Tensor;
    pub fn copy(&self) -> Tensor;
    pub fn from_slice2<T, U>(v: &[U]) -> Tensor;
    pub fn to_mkldnn(&self) -> Tensor;
}

// From traits
impl<T: Element> From<&[T]> for Tensor { ... }
impl<T: Element> From<T> for Tensor { ... }

Import

use tch::{Tensor, Shape, Kind, Device};

I/O Contract

Method	Input	Output	Error Condition
view(s)	Any Shape type	Reshaped Tensor	Incompatible size
zero_pad1d	3D tensor, left/right padding	Padded 3D Tensor	Not 3D: TchError::Shape
zero_pad2d	4D tensor, left/right/top/bottom	Padded 4D Tensor	Not 4D: TchError::Shape
cross_entropy_for_logits	Logits tensor + targets	Scalar loss Tensor	Shape mismatch
accuracy_for_logits	Logits tensor + targets	Scalar accuracy Tensor	Shape mismatch
onehot(labels)	Integer tensor + label count	Float tensor with extra dim	None

Shape Type	Example	Resulting Shape
[i64; 2]	[3, 4]	[3, 4]
(i64, i64)	(3, 4)	[3, 4]
i64	12	[12]
()	()	[] (scalar)
&[i64]	&[3, 4, 5]	[3, 4, 5]

Usage Examples

use tch::{Tensor, Kind, Device};

// Using Shape trait with view
let t = Tensor::randn([2, 3, 4], (Kind::Float, Device::Cpu));
let reshaped = t.view([6, 4]);       // [i64; 2]
let reshaped = t.view((2, 12));      // (i64, i64)
let reshaped = t.view(24i64);        // single i64

// Zero padding
let t3d = Tensor::randn([1, 3, 10], (Kind::Float, Device::Cpu));
let padded = t3d.zero_pad1d(2, 2);  // shape: [1, 3, 14]

let t4d = Tensor::randn([1, 3, 8, 8], (Kind::Float, Device::Cpu));
let padded = t4d.zero_pad2d(1, 1, 1, 1);  // shape: [1, 3, 10, 10]

// Loss computation
let logits = Tensor::randn([32, 10], (Kind::Float, Device::Cpu));
let targets = Tensor::randint(10, [32], (Kind::Int64, Device::Cpu));
let loss = logits.cross_entropy_for_logits(&targets);
let acc = logits.accuracy_for_logits(&targets);

// One-hot encoding
let labels = Tensor::from_slice(&[0i64, 2, 1]);
let onehot = labels.onehot(3);  // shape: [3, 3]

// Create from slice and scalar
let t = Tensor::from(&[1.0f32, 2.0, 3.0][..]);
let s = Tensor::from(42.0f64);

Related Pages

• Principle:LaurentMazare_Tch_rs_Tensor_Shape_Abstraction