Static library 2.4

CHANGELOG.md

@@ -7,6 +7,18 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## Unreleased
 ### Changed
 
+## v0.16.0
+### Changed
+- PyTorch v2.4 support
+
+## v0.16.0
+### Changed
+- PyTorch v2.3 support
+
+## v0.15.0
+### Changed
+- PyTorch v2.2 support
+
 ## v0.14.0
 ### Changed
 - PyTorch v2.1 support

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "tch"
-version = "0.14.0"
+version = "0.17.0"
 authors = ["Laurent Mazare <[email protected]>"]
 edition = "2021"
 build = "build.rs"
@@ -23,7 +23,7 @@ ndarray = "0.15"
 num-traits = "0.2.15"
 rand = "0.8"
 thiserror = "1"
-torch-sys = { version = "0.14.0", path = "torch-sys" }
+torch-sys = { version = "0.17.0", path = "torch-sys" }
 zip = { version = "0.6", default-features = false, features = [
     # Any other features, in particular `deflate-<backend>` features
     # may be added by other crates and override the default backend.

README.md

@@ -18,7 +18,7 @@ The code generation part for the C api on top of libtorch comes from
 
 ## Getting Started
 
-This crate requires the C++ PyTorch library (libtorch) in version *v2.1.0* to be available on
+This crate requires the C++ PyTorch library (libtorch) in version *v2.4.0* to be available on
 your system. You can either:
 
 - Use the system-wide libtorch installation (default).
@@ -54,9 +54,9 @@ export LIBTORCH=/path/to/libtorch
 The header files location can also be specified separately from the shared library via
 the following:
 ```bash
-# LIBTORCH_INCLUDE must contains `include` directory.
+# LIBTORCH_INCLUDE must contain `include` directory.
 export LIBTORCH_INCLUDE=/path/to/libtorch/
-# LIBTORCH_LIB must contains `lib` directory.
+# LIBTORCH_LIB must contain `lib` directory.
 export LIBTORCH_LIB=/path/to/libtorch/
 ```
 - For Windows users, assuming that `X:\path\to\libtorch` is the unzipped libtorch directory.
@@ -85,7 +85,7 @@ seem to include `libtorch.a` by default so this would have to be compiled
 manually, e.g. via the following:
 
 ```bash
-git clone -b v2.1.0 --recurse-submodule https://github.com/pytorch/pytorch.git pytorch-static --depth 1
+git clone -b v2.4.0 --recurse-submodule https://github.com/pytorch/pytorch.git pytorch-static --depth 1
 cd pytorch-static
 USE_CUDA=OFF BUILD_SHARED_LIBS=OFF python setup.py build
 # export LIBTORCH to point at the build directory in pytorch-static.

examples/min-gpt/main.rs

@@ -80,7 +80,7 @@ fn causal_self_attention(p: &nn::Path, cfg: Config) -> impl ModuleT {
         let q = xs.apply(&query).view(sizes).transpose(1, 2);
         let v = xs.apply(&value).view(sizes).transpose(1, 2);
         let att = q.matmul(&k.transpose(-2, -1)) * (1.0 / f64::sqrt(sizes[3] as f64));
-        let att = att.masked_fill(&mask.i((.., .., ..sz_t, ..sz_t)).eq(0.), std::f64::NEG_INFINITY);
+        let att = att.masked_fill(&mask.i((.., .., ..sz_t, ..sz_t)).eq(0.), f64::NEG_INFINITY);
         let att = att.softmax(-1, Kind::Float).dropout(cfg.attn_pdrop, train);
         let ys = att.matmul(&v).transpose(1, 2).contiguous().view([sz_b, sz_t, sz_c]);
         ys.apply(&proj).dropout(cfg.resid_pdrop, train)

examples/python-extension/Cargo.toml

@@ -17,7 +17,7 @@ name = "tch_ext"
 crate-type = ["cdylib"]
 
 [dependencies]
-pyo3 = { version = "0.18.3", features = ["extension-module"] }
-pyo3-tch = { path = "../../pyo3-tch", version = "0.14.0" }
-tch = { path = "../..", features = ["python-extension"], version = "0.14.0" }
-torch-sys = { path = "../../torch-sys", features = ["python-extension"], version = "0.14.0" }
+pyo3 = { version = "0.21", features = ["extension-module"] }
+pyo3-tch = { path = "../../pyo3-tch", version = "0.17.0" }
+tch = { path = "../..", features = ["python-extension"], version = "0.17.0" }
+torch-sys = { path = "../../torch-sys", features = ["python-extension"], version = "0.17.0" }

examples/python-extension/src/lib.rs

@@ -3,15 +3,15 @@ use pyo3_tch::{wrap_tch_err, PyTensor};
 
 #[pyfunction]
 fn add_one(tensor: PyTensor) -> PyResult<PyTensor> {
-    let tensor = tensor.f_add_scalar(1.0).map_err(wrap_tch_err)?;
+    let tensor = tensor.f_add_scalar(1.0, 1.0).map_err(wrap_tch_err)?;
     Ok(PyTensor(tensor))
 }
 
 /// A Python module implemented in Rust using tch to manipulate PyTorch
 /// objects.
 #[pymodule]
-fn tch_ext(py: Python<'_>, m: &PyModule) -> PyResult<()> {
-    py.import("torch")?;
+fn tch_ext(py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
+    py.import_bound("torch")?;
     m.add_function(wrap_pyfunction!(add_one, m)?)?;
     Ok(())
 }

examples/stable-diffusion/main.rs

@@ -36,7 +36,6 @@
 //
 //   cargo run --release --example tensor-tools cp ./data/vae.npz ./data/vae.ot
 //   cargo run --release --example tensor-tools cp ./data/unet.npz ./data/unet.ot
-///
 // TODO: fix tensor_tools so that it works properly there.
 // TODO: Split this file, probably in a way similar to huggingface/diffusers.
 use std::collections::{HashMap, HashSet};

examples/yolo/darknet.rs

@@ -17,7 +17,7 @@ struct Block {
 
 impl Block {
     fn get(&self, key: &str) -> Result<&str> {
-        match self.parameters.get(&key.to_string()) {
+        match self.parameters.get(key) {
             None => bail!("cannot find {} in {}", key, self.block_type),
             Some(value) => Ok(value),
         }
@@ -32,7 +32,7 @@ pub struct Darknet {
 
 impl Darknet {
     fn get(&self, key: &str) -> Result<&str> {
-        match self.parameters.get(&key.to_string()) {
+        match self.parameters.get(key) {
             None => bail!("cannot find {} in net parameters", key),
             Some(value) => Ok(value),
         }
@@ -199,7 +199,7 @@ where
     slice.copy_(&src)
 }
 
-fn detect(xs: &Tensor, image_height: i64, classes: i64, anchors: &Vec<(i64, i64)>) -> Tensor {
+fn detect(xs: &Tensor, image_height: i64, classes: i64, anchors: &[(i64, i64)]) -> Tensor {
     let (bsize, _channels, height, _width) = xs.size4().unwrap();
     let stride = image_height / height;
     let grid_size = image_height / stride;

gen/gen.ml

@@ -32,6 +32,7 @@ let excluded_functions =
     ; "_cummax_helper"
     ; "retain_grad"
     ; "_validate_sparse_coo_tensor_args"
+    ; "_sparse_semi_structured_addmm"
     ; "_backward"
     ; "size"
     ; "stride"
@@ -101,7 +102,9 @@ let excluded_prefixes =
   ; "_amp_foreach"
   ; "_nested_tensor"
   ; "_fused_adam"
+  ; "_fused_adagrad"
   ; "sym_"
+  ; "_fused_sgd"
   ]
 
 let excluded_suffixes = [ "_forward"; "_forward_out" ]
@@ -176,6 +179,7 @@ module Func = struct
     | "at::tensoroptions" -> Some TensorOptions
     | "at::intarrayref" -> Some (if is_nullable then IntListOption else IntList)
     | "at::arrayref<double>" -> Some DoubleList
+    | "const c10::list<::std::optional<at::tensor>> &"
     | "const c10::list<c10::optional<at::tensor>> &" -> Some TensorOptList
     | "const at::itensorlistref &" | "at::tensorlist" -> Some TensorList
     | "at::device" -> Some Device
@@ -598,7 +602,7 @@ let write_cpp funcs filename =
       let pc s = p out_cpp s in
       let ph s = p out_h s in
       pc "// THIS FILE IS AUTOMATICALLY GENERATED, DO NOT EDIT BY HAND!";
-      pc "#include \"%s.h\"" (Caml.Filename.basename filename);
+      pc "#include \"%s.h\"" (Stdlib.Filename.basename filename);
       pc "";
       ph "// THIS FILE IS AUTOMATICALLY GENERATED, DO NOT EDIT BY HAND!";
       ph "#include \"torch_api.h\"";
@@ -887,7 +891,7 @@ let run
 
 let () =
   run
-    ~yaml_filename:"third_party/pytorch/Declarations-v2.1.0.yaml"
+    ~yaml_filename:"third_party/pytorch/Declarations-v2.4.0.yaml"
     ~cpp_filename:"torch-sys/libtch/torch_api_generated"
     ~ffi_filename:"torch-sys/src/c_generated.rs"
     ~wrapper_filename:"src/wrappers/tensor_generated.rs"

pyo3-tch/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "pyo3-tch"
-version = "0.14.0"
+version = "0.17.0"
 authors = ["Laurent Mazare <[email protected]>"]
 edition = "2021"
 build = "build.rs"
@@ -12,6 +12,6 @@ categories = ["science"]
 license = "MIT/Apache-2.0"
 
 [dependencies]
-tch = { path = "..", features = ["python-extension"], version = "0.14.0" }
-torch-sys = { path = "../torch-sys", features = ["python-extension"], version = "0.14.0" }
-pyo3 = { version = "0.18.3", features = ["extension-module"] }
+tch = { path = "..", features = ["python-extension"], version = "0.17.0" }
+torch-sys = { path = "../torch-sys", features = ["python-extension"], version = "0.17.0" }
+pyo3 = { version = "0.21", features = ["extension-module"] }

pyo3-tch/src/lib.rs

@@ -1,8 +1,5 @@
+use pyo3::exceptions::{PyTypeError, PyValueError};
 use pyo3::prelude::*;
-use pyo3::{
-    exceptions::{PyTypeError, PyValueError},
-    AsPyPointer,
-};
 pub use tch;
 pub use torch_sys;
 

src/nn/init.rs

@@ -107,7 +107,7 @@ pub fn f_init(i: Init, dims: &[i64], device: Device) -> Result<Tensor, TchError>
             // Optimize the case for which a single C++ code can be done.
             if cst == 0. {
                 Tensor::f_zeros(dims, (Kind::Float, device))
-            } else if (cst - 1.).abs() <= std::f64::EPSILON {
+            } else if (cst - 1.).abs() <= f64::EPSILON {
                 Tensor::f_ones(dims, (Kind::Float, device))
             } else {
                 Tensor::f_ones(dims, (Kind::Float, device)).map(|t| t * cst)
@@ -117,7 +117,7 @@ pub fn f_init(i: Init, dims: &[i64], device: Device) -> Result<Tensor, TchError>
             Tensor::f_zeros(dims, (Kind::Float, device))?.f_uniform_(lo, up)
         }
         Init::Randn { mean, stdev } => {
-            if mean == 0. && (stdev - 1.).abs() <= std::f64::EPSILON {
+            if mean == 0. && (stdev - 1.).abs() <= f64::EPSILON {
                 Tensor::f_randn(dims, (Kind::Float, device))
             } else {
                 Tensor::f_randn(dims, (Kind::Float, device)).map(|t| t * stdev + mean)

src/nn/var_store.rs

@@ -151,6 +151,11 @@ impl VarStore {
     ///
     /// Weight values for all the tensors currently stored in the
     /// var-store are saved in the given file.
+    ///
+    /// If the given path ends with the suffix `.safetensors`, the file will
+    /// be saved in safetensors format. Otherwise, libtorch C++ module format
+    /// will be used. Note that saving in pickle format (`.pt` extension) is
+    /// not supported by the C++ API of Torch.
     pub fn save<T: AsRef<std::path::Path>>(&self, path: T) -> Result<(), TchError> {
         let variables = self.variables_.lock().unwrap();
         let named_tensors = variables.named_variables.iter().collect::<Vec<_>>();
@@ -216,6 +221,11 @@ impl VarStore {
     /// var-store are loaded from the given file. Note that the set of
     /// variables stored in the var-store is not changed, only the values
     /// for these tensors are modified.
+    ///
+    /// The format of the file is deduced from the file extension:
+    /// - `.safetensors`: The file is assumed to be in safetensors format.
+    /// - `.bin` or `.pt`: The file is assumed to be in pickle format.
+    /// - Otherwise, the file is assumed to be in libtorch C++ module format.
     pub fn load<T: AsRef<std::path::Path>>(&mut self, path: T) -> Result<(), TchError> {
         if self.device != Device::Mps {
             self.load_internal(path)