Support for instantiating tensors from num_complex::Complex values.

Cargo.toml

@@ -33,6 +33,7 @@ image = { version = "0.24.5", optional = true }
 clap = { version = "4.2.4", features = ["derive"], optional = true }
 serde_json = { version = "1.0.96", optional = true }
 memmap2 = { version = "0.6.1", optional = true }
+num-complex = { version = "0.4.6", optional = true }
 
 [dev-dependencies]
 anyhow = "^1.0.60"
@@ -50,6 +51,7 @@ python-extension = ["torch-sys/python-extension"]
 rl-python = ["cpython"]
 doc-only = ["torch-sys/doc-only"]
 cuda-tests = []
+complex = ["dep:num-complex"]
 
 [package.metadata.docs.rs]
 features = [ "doc-only" ]

examples/cplx_fft.rs

@@ -0,0 +1,19 @@
+#[cfg(feature = "complex")]
+fn main() {
+    use std::f64::consts::PI;
+    use tch::*;
+
+    let len = 32;
+    let freq = 0.2;
+
+    let n = Tensor::arange(len, (Kind::Float, Device::Cpu));
+    let theta = &n * (2.0 * PI * freq);
+    let real = theta.cos();
+    let imag = theta.sin();
+    let sinusoid = Tensor::complex(&real, &imag);
+    let fft_bins = sinusoid.fft_fft(len, -1, "backward");
+    let psd = fft_bins.abs().pow_tensor_scalar(1.0) / len as f64;
+    let psd: Vec<f32> = psd.try_into().unwrap();
+    let max_idx = psd.iter().enumerate().max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()).map(|(idx, _)| idx).unwrap();
+    println!("Frequency: approx = {}, actual = {}", max_idx as f32 / len as f32, freq);
+}

src/wrappers/kind.rs

@@ -174,3 +174,15 @@ unsafe impl Element for bool {
     const KIND: Kind = Kind::Bool;
     const ZERO: Self = false;
 }
+
+#[cfg(feature = "complex")]
+unsafe impl Element for num_complex::Complex<f32> {
+    const KIND: Kind = Kind::ComplexFloat;
+    const ZERO: Self = num_complex::Complex::<f32>::ZERO;
+}
+
+#[cfg(feature = "complex")]
+unsafe impl Element for num_complex::Complex<f64> {
+    const KIND: Kind = Kind::ComplexDouble;
+    const ZERO: Self = num_complex::Complex::<f64>::ZERO;
+}

tests/tensor_tests.rs

@@ -3,7 +3,7 @@ use anyhow::Result;
 use half::f16;
 use std::convert::{TryFrom, TryInto};
 use std::f32;
-use tch::{Device, TchError, Tensor};
+use tch::{Device, Kind, TchError, Tensor};
 
 mod test_utils;
 use test_utils::*;
@@ -215,6 +215,17 @@ fn into_ndarray_i64() {
     assert_eq!(vec_i64_from(&tensor).as_slice(), nd.as_slice().unwrap());
 }
 
+#[cfg(feature = "complex")]
+#[test]
+fn into_ndarray_cplx() {
+    use num_complex::{c32, Complex};
+    let tensor = Tensor::from_slice(&[c32(0., 0.), c32(1. , 0.), c32(0., 1.), c32(1.,1.)]).reshape(&[2, 2]);
+    let nd: ndarray::ArrayD<Complex<f32>> = (&tensor).try_into().unwrap();
+    assert_eq!(nd.shape(), [2, 2]);
+    assert_eq!(vec_cplx_from(&tensor).as_slice(), nd.as_slice().unwrap());
+}
+
+
 #[test]
 fn from_ndarray_f64() {
     let nd = ndarray::arr2(&[[1f64, 2.], [3., 4.]]);
@@ -236,6 +247,16 @@ fn from_ndarray_bool() {
     assert_eq!(vec_bool_from(&tensor).as_slice(), nd.as_slice().unwrap());
 }
 
+#[cfg(feature = "complex")]
+#[test]
+fn from_ndarray_cplx() {
+    use num_complex::c32;
+    let nd = ndarray::arr2(&[[c32(0., 0.), c32(1. , 0.)], [c32(0., 1.), c32(1.,1.)]]);
+    let tensor = Tensor::try_from(nd.clone()).unwrap();
+    assert_eq!(tensor.dim(), 2);
+    assert_eq!(vec_cplx_from(&tensor).as_slice(), nd.as_slice().unwrap());
+}
+
 #[test]
 fn from_primitive() -> Result<()> {
     assert_eq!(vec_i32_from(&Tensor::try_from(1_i32)?), vec![1]);
@@ -506,3 +527,31 @@ fn fp8_tensor() {
         assert_eq!(t.kind().elt_size_in_bytes(), 1);
     }
 }
+
+#[cfg(feature = "complex")]
+#[test]
+fn complex_tensor() {
+    use num_complex::Complex;
+
+    let input = vec![
+        Complex::<f32>::new(0., 0.),
+        Complex::<f32>::new(1., 0.),
+        Complex::<f32>::new(1., 1.),
+        Complex::<f32>::new(0., 1.),
+    ];
+
+    let t = Tensor::from_slice(&input);
+    assert_eq!(t.kind(), Kind::ComplexFloat);
+    assert_eq!(t.kind().elt_size_in_bytes(), 8);
+    assert!(t.is_complex());
+
+    let c = t.conj_physical();
+    let back: Vec<Complex<f32>> = c.try_into().unwrap();
+    let expected = input.iter().map(Complex::conj).collect::<Vec<_>>();
+    assert_eq!(expected, back);
+
+    let t = Tensor::from_slice(&[Complex::<f64>::new(3., -7.)]);
+    assert_eq!(t.kind(), Kind::ComplexDouble);
+    assert_eq!(t.kind().elt_size_in_bytes(), 16);
+
+}

tests/test_utils.rs

@@ -42,3 +42,9 @@ pub fn vec_i32_from(t: &Tensor) -> Vec<i32> {
 pub fn vec_bool_from(t: &Tensor) -> Vec<bool> {
     from::<Vec<bool>>(&t.reshape(-1))
 }
+
+#[cfg(feature = "complex")]
+#[allow(dead_code)]
+pub fn vec_cplx_from(t: &Tensor) -> Vec<num_complex::Complex<f32>> {
+    from::<Vec<num_complex::Complex<f32>>>(&t.reshape(-1))
+}