Implementation:LaurentMazare Tch rs Crate Root
| Knowledge Sources | |
|---|---|
| Domains | Crate Organization, Module System |
| Last Updated | 2026-02-08 00:00 GMT |
Overview
The crate root src/lib.rs is the main entry point for the tch crate, defining all public module declarations, re-exports, and the crate-wide Result type alias.
Description
This file establishes the top-level public API surface of the tch crate. It uses #[macro_use] extern crate lazy_static to enable lazy static initialization throughout the crate. The file declares five primary modules: data (public), error (private with re-export), wrappers (crate-private), tensor (private with re-exports), nn (public), and vision (public).
Key types are selectively re-exported from internal modules to provide a flat, ergonomic import experience. The TchError enum and a Result<T> type alias are established at the crate root. Wrapper types like Device, Cuda, Kind, Layout, CModule, IValue, Scalar, and COptimizer are re-exported from wrappers. Tensor-related types including Tensor, Shape, NoGradGuard, Reduction, and index operators are re-exported from the tensor module. The python module is conditionally compiled behind the python-extension feature flag.
Usage
Users import from tch directly without needing to navigate internal module paths. For example, use tch::Tensor or use tch::nn provides access to the full library API.
Code Reference
Source Location
- Repository: LaurentMazare_Tch_rs
- File: src/lib.rs
Signature
// Error type and result alias
pub use error::TchError;
pub type Result<T> = std::result::Result<T, error::TchError>;
// Wrapper re-exports
pub use wrappers::device::{Cuda, Device};
pub use wrappers::jit::{self, CModule, IValue, TrainableCModule};
pub use wrappers::kind::{self, Kind};
pub use wrappers::layout::Layout;
pub use wrappers::optimizer::COptimizer;
pub use wrappers::scalar::Scalar;
pub use wrappers::utils;
pub use wrappers::{
get_num_interop_threads, get_num_threads, manual_seed,
set_num_interop_threads, set_num_threads, QEngine,
};
// Tensor re-exports
pub use tensor::{
autocast, display, index, no_grad, no_grad_guard, with_grad,
IndexOp, NewAxis, NoGradGuard, Reduction, Shape, Tensor, TensorIndexer,
};
// Conditional module
#[cfg(feature = "python-extension")]
pub use wrappers::python;
// Public modules
pub mod data;
pub mod nn;
pub mod vision;
Import
use tch::{Tensor, Device, Kind, nn, vision};
use tch::{TchError, Result};
I/O Contract
| Module | Visibility | Key Re-exports |
|---|---|---|
| data | Public | Dataset utilities |
| error | Private (re-exported) | TchError, Result<T> |
| wrappers | Crate-private | Device, Cuda, Kind, Layout, CModule, IValue, TrainableCModule, COptimizer, Scalar, QEngine, threading functions, manual_seed |
| tensor | Private (re-exported) | Tensor, Shape, NoGradGuard, Reduction, IndexOp, NewAxis, TensorIndexer, autocast, no_grad, with_grad |
| nn | Public | Neural network layers and optimizers |
| vision | Public | Vision dataset loaders and image models |
| Feature Flag | Effect |
|---|---|
| python-extension | Enables pub use wrappers::python for Python tensor interop |
Usage Examples
// The crate root enables flat imports
use tch::{Tensor, Device, Kind};
use tch::nn::{self, Module, OptimizerConfig};
// Create a tensor using re-exported types
let t = Tensor::randn([3, 4], (Kind::Float, Device::Cpu));
// Use the Result type alias
fn example() -> tch::Result<Tensor> {
Ok(Tensor::zeros([2, 2], (Kind::Float, Device::Cpu)))
}