[Python] How to safely store tch::Tensor in a struct that requires Sync?

I'm developing a Python library in Rust using PyO3 0.24.2 and tch-rs 0.20.0, and I need to store tch::Tensor objects in a Rust struct that is exposed as a Python class via #[pyclass]. However, I'm hitting a compile-time error because tch::Tensor is not Sync, which is required by PyO3 from version 0.23.0 for #[pyclass] types due to Python's threading model.

Here's a minimal example of my code:

use pyo3:relude::*;
use tch::Tensor;

struct ClassWithTensor {
tensor: tch::Tensor
}

#[pyclass]
struct PyClassWithTensor(ClassWithTensor);

#[pymodule]
fn init_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
m.add_class::<PyClassWithTensor>()?;
Ok(())
}


This fails to compile with the following error:

error[E0277]: `*mut torch_sys::C_tensor` cannot be shared between threads safely
--> src/lib.rs:9:8
|
9 | struct PyClassWithTensor(ClassWithTensor);
| ^^^^^^^^^^^^^^^^^ `*mut torch_sys::C_tensor` cannot be shared between threads safely
|
= help: within `PyClassWithTensor`, the trait `Sync` is not implemented for `*mut torch_sys::C_tensor`
note: required because it appears within the type `tch::Tensor`
--> /home/antoni/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/tch-0.20.0/src/wrappers/tensor.rs:18:12
|
18 | pub struct Tensor {
| ^^^^^^
note: required because it appears within the type `ClassWithTensor`
--> src/lib.rs:4:8
|
4 | struct ClassWithTensor {
| ^^^^^^^^^^^^^^^
note: required because it appears within the type `PyClassWithTensor`
--> src/lib.rs:9:8
|
9 | struct PyClassWithTensor(ClassWithTensor);
| ^^^^^^^^^^^^^^^^^
note: required by a bound in `assert_pyclass_sync`
--> /home/antoni/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/pyo3-0.24.2/src/impl_/pyclass/assertions.rs:12:22
|
10 | pub const fn assert_pyclass_sync<T>()
| ------------------- required by a bound in this function
11 | where
12 | T: PyClassSync + Sync,
| ^^^^ required by this bound in `assert_pyclass_sync`


I found a related discussion in tch-rs Issue #460, where a user proposed wrapping tch::Tensor in a TensorWrapper struct with Arc and implementing Send and Sync unsafely:

struct TensorWrapper {
tensor: Arc<Tensor>
}
unsafe impl Send for TensorWrapper {}
unsafe impl Sync for TensorWrapper {}


However, the discussion suggests that tch::Tensor is not thread-safe enough to be Sync, and even Send might be risky (see also LaurentMazare/tch-rs#461).

The only example of a Python module written in Rust with tch-rs that I could find is tch-ext. However, this module uses older versions of PyO3 (0.18.3) and tch-rs (via pyo3-tch 0.15.0), and it does not compile with PyO3 0.24.2 and tch-rs 0.20.0. Additionally, tch-ext only demonstrates #[pyfunction] for tensor manipulation, not #[pyclass], so it doesn't address my use case of storing tch::Tensor in a struct.

I also noticed that PyTorch 2.7.0 supports Python 3.13t (free-threaded mode with GIL disabled) across platforms, as indicated in pytorch/pytorch#130249. This suggests that PyTorch's underlying C++ tensors (at::Tensor) are likely thread-safe in free-threaded mode. Maybe tch::Tensor can now be safely marked as Sync?

Questions:

1. How can I safely store tch::Tensor in a #[pyclass] struct for a Python library using PyO3 0.24.2 and tch-rs 0.20.0, ensuring compatibility with PyO3's Send + Sync requirements?
2. Are there alternative approaches to storing tch::Tensor in a #[pyclass] struct, such as wrapping it in synchronization primitives like Arc<Mutex>?
3. What risks should I be aware of when storing tch::Tensor in a #[pyclass] struct?

Thanks for any guidance or examples you can share!

Continue reading...