From 3bdebae1f6456a5b989f426006d153e9d832989f Mon Sep 17 00:00:00 2001
From: Ayoub Benaissa <ayoub.benaissa@zama.ai>
Date: Tue, 20 Jun 2023 09:07:03 +0100
Subject: [PATCH] fix(compiler): fix plaintext tensor shape when using crt

client parameter generation was removing last dimension of a tensor when using CRT including plaintext ones, while it should only be done for encrypted ones.
---
 .../compiler/lib/Support/ClientParametersGeneration.cpp    | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/compilers/concrete-compiler/compiler/lib/Support/ClientParametersGeneration.cpp b/compilers/concrete-compiler/compiler/lib/Support/ClientParametersGeneration.cpp
index 5ebb4408d..885829abd 100644
--- a/compilers/concrete-compiler/compiler/lib/Support/ClientParametersGeneration.cpp
+++ b/compilers/concrete-compiler/compiler/lib/Support/ClientParametersGeneration.cpp
@@ -210,9 +210,10 @@ generateGate(mlir::Type type, encodings::Encoding encoding,
         if (auto err = scalarGate.takeError()) {
           return std::move(err);
         }
-        if (maybeCrt.has_value()) {
-          // When using crt, the last dimension of the tensor is for the members
-          // of the decomposition. It should not be used.
+        if (maybeCrt.has_value() && scalarGate->isEncrypted()) {
+          // When using crt with encrypted tensors, the last dimension of the
+          // tensor is for the members of the decomposition. It should not be
+          // used.
           scalarGate->shape.dimensions =
               tensor.getShape().take_front(tensor.getShape().size() - 1).vec();
         } else {