From 15f9d2642fa0c6cabb7a1af7055c87fa1fb4a272 Mon Sep 17 00:00:00 2001
From: Guillaume Lagrange <lagrange.guillaume.1@gmail.com>
Date: Tue, 24 Sep 2024 15:53:47 -0400
Subject: [PATCH] Add note on handles/streams order

---
 crates/burn-fusion/src/client/mutex.rs | 2 +-
 crates/burn-jit/src/fusion/base.rs     | 1 +
 2 files changed, 2 insertions(+), 1 deletion(-)

diff --git a/crates/burn-fusion/src/client/mutex.rs b/crates/burn-fusion/src/client/mutex.rs
index 0c99be132..bcc1a87e9 100644
--- a/crates/burn-fusion/src/client/mutex.rs
+++ b/crates/burn-fusion/src/client/mutex.rs
@@ -211,7 +211,7 @@ where
         core::mem::drop(server_other);
         core::mem::drop(server_current);
 
-        // Expected order [qtensor, scale, <offset>]
+        // NOTE: the expected order is known [qtensor, scale, <offset>]
         let offset = tensor.qparams.offset.map(|desc| {
             FusionTensor::new(
                 ids.pop().unwrap(),
diff --git a/crates/burn-jit/src/fusion/base.rs b/crates/burn-jit/src/fusion/base.rs
index 36e06c628..e0dcde3bc 100644
--- a/crates/burn-jit/src/fusion/base.rs
+++ b/crates/burn-jit/src/fusion/base.rs
@@ -82,6 +82,7 @@ impl<R: JitRuntime, F: FloatElement, I: IntElement> ReprBackend for JitBackend<R
         scheme: QuantizationScheme,
     ) -> burn_tensor::ops::QuantizedTensor<Self> {
         match handles.len() {
+            // NOTE: the order of the handles is known [qtensor, scale, <offset>]
             3 => {
                 let mut handles = handles;
                 let offset = handles.pop().unwrap();