feat: update readme

Author: dancixx

README.md

@@ -4,7 +4,7 @@
 
 # iTransformer: Rust Implementation
 
-An **iTransformer** implementation in **Rust**, inspired by the [lucidrains iTransformer repository](https://github.com/lucidrains/iTransformer), and based on the original research and implementation from [Tsinghua University's iTransformer repository](https://github.com/thuml/iTransformer).
+An **iTransformer** implementation in **Rust** using [Hugging Face Candle](https://github.com/huggingface/candle), inspired by the [lucidrains iTransformer repository](https://github.com/lucidrains/iTransformer), and based on the original research and implementation from [Tsinghua University's iTransformer repository](https://github.com/thuml/iTransformer).
 
 ## 📚 **What is iTransformer?**
 
@@ -16,7 +16,8 @@ iTransformer introduces an **inverted Transformer architecture** designed for **
 - **Flexible Prediction Lengths:** Supports predictions at multiple horizons (e.g., 12, 24, 36, 48 steps ahead).
 - **Scalability:** Handles hundreds of variates with efficiency.
 - **Zero-Shot Generalization:** Train on partial variates and generalize to unseen variates.
-- **Efficient Attention Mechanisms:** Compatible with advanced techniques like FlashAttention.
+- **Multiple Model Variants:** ITransformer, ITransformer2D, and ITransformerFFT.
+- **Hardware Acceleration:** Metal (macOS) and CUDA support via Candle.
 
 ### 🛠️ **Architecture Overview**
 
@@ -41,42 +42,154 @@ To get started, ensure you have Rust and Cargo installed. Then:
 
 ```bash
 # Add iTransformer to your project dependencies
-cargo add itransformer-rs
+cargo add itransformer
 ```
 
+For GPU acceleration:
+```bash
+# macOS Metal support
+cargo add itransformer --features metal
+
+# CUDA support
+cargo add itransformer --features cuda
+```
+
+## 🏗️ **Model Variants**
+
+This library provides three model variants:
+
+| Model | Description |
+|-------|-------------|
+| `ITransformer` | Base inverted transformer for multivariate time series forecasting |
+| `ITransformer2D` | Extended variant with granular time attention via `num_time_tokens` parameter |
+| `ITransformerFFT` | Variant with additional Fourier tokens prepended to the attention sequence |
+
 ## 📝 **Usage**
 
-```rust
-use tch::{Device, Tensor, nn::VarStore, Kind};
+### ITransformer (Base Model)
 
-fn main() -> Result<(), Box<dyn std::error::Error>> {
-    let vs = VarStore::new(Device::Cpu);
+```rust
+use candle_core::{Device, DType, Tensor};
+use candle_nn::{VarBuilder, VarMap};
+use itransformer::ITransformer;
+
+fn main() -> candle_core::Result<()> {
+    let device = Device::Cpu;
+    let varmap = VarMap::new();
+    let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
+    
     let model = ITransformer::new(
-        &(vs.root() / "itransformer"),
-        137, // num_variates
-        96,  // lookback_len
-        6,   // depth
-        256, // dim
-        Some(1),       // num_tokens_per_variate
-        vec![12, 24, 36, 48], // pred_length
-        Some(64),     // dim_head
-        Some(8),      // heads
-        None,         // attn_drop_p
-        None,         // ff_mult
-        None,         // ff_drop_p
-        None,         // num_mem_tokens
-        Some(true),   // use_reversible_instance_norm
-        None,         // reversible_instance_norm_affine
-        false,        // flash_attn
-        &Device::Cpu,
+        vb,
+        137,                      // num_variates
+        96,                       // lookback_len
+        6,                        // depth
+        256,                      // dim
+        Some(1),                  // num_tokens_per_variate
+        vec![12, 24, 36, 48],     // pred_length
+        Some(64),                 // dim_head
+        Some(8),                  // heads
+        None,                     // attn_drop_p
+        None,                     // ff_mult
+        None,                     // ff_drop_p
+        None,                     // num_mem_tokens
+        Some(true),               // use_reversible_instance_norm
+        None,                     // reversible_instance_norm_affine
+        false,                    // flash_attn
+        &device,
     )?;
-    let time_series = Tensor::randn([2, 96, 137], (Kind::Float, Device::Cpu));
-    let preds = model.forward(&time_series, None, false);
+    
+    let time_series = Tensor::randn(0f32, 1f32, (2, 96, 137), &device)?;
+    let preds = model.forward(&time_series, None, false)?;
     println!("{:?}", preds);
     Ok(())
 }
 ```
 
+### ITransformer2D (Time Token Variant)
+
+```rust
+use candle_core::{Device, DType, Tensor};
+use candle_nn::{VarBuilder, VarMap};
+use itransformer::ITransformer2D;
+
+fn main() -> candle_core::Result<()> {
+    let device = Device::Cpu;
+    let varmap = VarMap::new();
+    let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
+    
+    let model = ITransformer2D::new(
+        vb,
+        137,                      // num_variates
+        96,                       // lookback_len
+        4,                        // depth
+        256,                      // dim
+        8,                        // num_time_tokens (lookback_len must be divisible by this)
+        vec![12, 24],             // pred_length
+        Some(32),                 // dim_head
+        Some(4),                  // heads
+        None, None, None, None, None, None,
+        false,                    // flash_attn
+        &device,
+    )?;
+    
+    let time_series = Tensor::randn(0f32, 1f32, (2, 96, 137), &device)?;
+    let preds = model.forward(&time_series, None, false)?;
+    println!("{:?}", preds);
+    Ok(())
+}
+```
+
+### ITransformerFFT (Fourier Token Variant)
+
+```rust
+use candle_core::{Device, DType, Tensor};
+use candle_nn::{VarBuilder, VarMap};
+use itransformer::ITransformerFFT;
+
+fn main() -> candle_core::Result<()> {
+    let device = Device::Cpu;
+    let varmap = VarMap::new();
+    let vb = VarBuilder::from_varmap(&varmap, DType::F32, &device);
+    
+    let model = ITransformerFFT::new(
+        vb,
+        137,                      // num_variates
+        96,                       // lookback_len
+        4,                        // depth
+        256,                      // dim
+        Some(1),                  // num_tokens_per_variate
+        4,                        // num_fft_tokens
+        vec![12, 24],             // pred_length
+        Some(32),                 // dim_head
+        Some(4),                  // heads
+        None, None, None, None, None, None,
+        false,                    // flash_attn
+        &device,
+    )?;
+    
+    let time_series = Tensor::randn(0f32, 1f32, (2, 96, 137), &device)?;
+    let preds = model.forward(&time_series, None, false)?;
+    println!("{:?}", preds);
+    Ok(())
+}
+```
+
+## 📁 **Project Structure**
+
+```
+src/
+├── attend.rs           # Attention computation
+├── attention.rs        # ToQKV, ToValueResidualMix, ToVGates, ToOut, Attention
+├── feedforward.rs      # GEGLU, FeedForward
+├── itransformer.rs     # Base ITransformer
+├── itransformer2d.rs   # ITransformer2D with time tokens
+├── itransformer_fft.rs # ITransformerFFT with Fourier tokens
+├── lib.rs              # Module declarations and re-exports
+├── mlp_in.rs           # Input projection layer
+├── pred_head.rs        # Prediction heads
+├── revin.rs            # Reversible Instance Normalization
+```
+
 
 ## 📖 **References**