From 1fc6099fdebc0bb186c09b73907e1dc973772f0c Mon Sep 17 00:00:00 2001 From: "Vishal S. Pandey" Date: Sun, 12 Apr 2026 21:44:55 +0530 Subject: [PATCH] Cast outputs to input dtypes; add result CSVs Author: @Akshay594 Ensure _combine returns tensors with the same dtypes as its inputs to avoid unintended dtype promotion during matrix operations (cast matmul results to F_l.dtype and G_l.dtype). Add `# type: ignore` to optional Triton imports. Add placeholder result CSVs retrieval_VLACombined_256.csv and retrieval_VLAKVExploding_256.csv (one includes a header row). --- results/retrieval_VLACombined_256.csv | 1 + results/retrieval_VLAKVExploding_256.csv | 0 src/models/attention/fast_vla.py | 6 +++--- 3 files changed, 4 insertions(+), 3 deletions(-) create mode 100644 results/retrieval_VLACombined_256.csv create mode 100644 results/retrieval_VLAKVExploding_256.csv diff --git a/results/retrieval_VLACombined_256.csv b/results/retrieval_VLACombined_256.csv new file mode 100644 index 0000000..2f06fcc --- /dev/null +++ b/results/retrieval_VLACombined_256.csv @@ -0,0 +1 @@ +epoch,train_loss,train_acc,test_acc,epoch_time_sec diff --git a/results/retrieval_VLAKVExploding_256.csv b/results/retrieval_VLAKVExploding_256.csv new file mode 100644 index 0000000..e69de29 diff --git a/src/models/attention/fast_vla.py b/src/models/attention/fast_vla.py index 70b3b5c..54ad7ee 100644 --- a/src/models/attention/fast_vla.py +++ b/src/models/attention/fast_vla.py @@ -6,15 +6,15 @@ # Try importing Triton try: - import triton - import triton.language as tl + import triton # type: ignore + import triton.language as tl # type: ignore HAS_TRITON = True except ImportError: HAS_TRITON = False def _combine(F_l, G_l, F_r, G_r): """Associative operator for the (F, G) recurrence.""" - return torch.matmul(F_r, F_l), torch.matmul(F_r, G_l) + G_r + return torch.matmul(F_r, F_l).to(F_l.dtype), (torch.matmul(F_r, G_l) + G_r).to(G_l.dtype) def parallel_scan(Fs, Gs): """