diff --git a/include/aidge/data/DataFormat.hpp b/include/aidge/data/DataFormat.hpp
index 4df854a1ea880e3e7def7ff6aee798ab95d30254..9b929b14ad373f37f488df80d9ed7823d3da1b20 100644
--- a/include/aidge/data/DataFormat.hpp
+++ b/include/aidge/data/DataFormat.hpp
@@ -83,6 +83,25 @@ using DataFormatTranspose = std::array<std::size_t, 5>;
*/
DataFormatTranspose getPermutationMapping(const DataFormat& src, const DataFormat& dst);
+/**
+ * @brief Determine the output data format after applying a transpose operation.
+ *
+ * This function computes the new data format after transposing a tensor.
+ * It follows these steps:
+ * 1. Retrieves the permutation mapping from the default format (NCHW) to the input format.
+ * 2. Applies the given output dimensions order to compute the new permutation.
+ * 3. Determines the number of dimensions associated with the input data format.
+ * 4. Identifies the corresponding data format for the new permutation.
+ *
+ * @param inputDataFormat The data format of the input tensor.
+ * @param outputDimsOrder The permutation applied to the input tensor's dimensions.
+ * @return DataFormat The computed output data format after applying the permutation.
+ *
+ * @throws std::runtime_error if the provided permutation is invalid for the given format.
+ * @pre `outputDimsOrder` must be a valid permutation of the format's dimensions.
+ */
+DataFormat getTransposeOutputDataFormat(const DataFormat inputDataFormat, const std::array<std::size_t, 5> outputDimsOrder);
+
} // namespace Aidge
namespace {
diff --git a/src/data/DataFormat.cpp b/src/data/DataFormat.cpp
index 8b7460b3d42de2e2ddb60a6332405141f87b297c..25621e1af0f37b340cc008efdcc407dd3028e0e4 100644
--- a/src/data/DataFormat.cpp
+++ b/src/data/DataFormat.cpp
@@ -51,6 +51,25 @@ static std::size_t getNbDimensions(DataFormat dformat) {
return nbDimensions[static_cast<std::size_t>(dformat)];
}
+/**
+ * @brief Get the DataFormat corresponding to a given permutation vector.
+ *
+ * @param perm The permutation vector.
+ * @return DataFormat The matching DataFormat, or DataFormat::Default if not found.
+ */
+static DataFormat getDataFormatFromPermutation(const DataFormatTranspose& perm, const unsigned int nbDims) {
+ for (size_t i = 0; i < NB_DFORMAT; ++i) {
+ DataFormat candidateFormat = static_cast<DataFormat>(i);
+ const DataFormatTranspose& candidatePerm = getPermutationFromNCHW(candidateFormat);
+
+ if (std::equal(candidatePerm.begin(), candidatePerm.begin() + nbDims, perm.begin())) {
+ return candidateFormat;
+ }
+ }
+ Log::warn("No matching DataFormat found for the given permutation : {}", perm);
+ return DataFormat::Default;
+}
+
DataFormatTranspose getPermutationMapping(const DataFormat& src, const DataFormat& dst) {
AIDGE_ASSERT((src != DataFormat::Any && dst != DataFormat::Any), "Permutation is not defined for DataFormat::Any");
if (src == DataFormat::Default || dst == DataFormat::Default || src == dst) {
@@ -78,4 +97,31 @@ DataFormatTranspose getPermutationMapping(const DataFormat& src, const DataForma
return src_to_dst;
}
+
+/**
+ * @brief Applies one permutation to another.
+ *
+ * Given an initial permutation `perm1` and another permutation `perm2`,
+ * this function returns a new permutation that results from applying `perm2` to `perm1`.
+ *
+ * @param perm1 The base permutation (original ordering → intermediate ordering).
+ * @param perm2 The permutation to apply (intermediate ordering → final ordering).
+ * @return DataFormatTranspose The resultng permutation (original ordering → final ordering).
+ */
+static DataFormatTranspose applyPermutation(const DataFormatTranspose& perm1, const DataFormatTranspose& perm2) {
+ DataFormatTranspose permuted{};
+ for (std::size_t i = 0; i < perm1.size(); ++i) {
+ permuted[i] = perm2[perm1[i]];
+ }
+ return permuted;
+}
+
+DataFormat getTransposeOutputDataFormat(const DataFormat inputDataFormat, const std::array<std::size_t, 5> outputDimsOrder){
+ const DataFormatTranspose permFromInput = getPermutationFromNCHW(inputDataFormat);
+ const DataFormatTranspose newPerm = applyPermutation(permFromInput, outputDimsOrder);
+
+ const unsigned int nbDims = getNbDimensions(inputDataFormat);
+
+ return getDataFormatFromPermutation(newPerm, nbDims);
+}
} // namespace Aidge
\ No newline at end of file
diff --git a/src/operator/Transpose.cpp b/src/operator/Transpose.cpp
index f9d612353a5fe8764419d6ac2f7fe1702f2a5df8..51c2d59605c2a6c3bb5c17b24663191c8e004bd0 100644
--- a/src/operator/Transpose.cpp
+++ b/src/operator/Transpose.cpp
@@ -77,7 +77,14 @@ bool Aidge::Transpose_Op::forwardDims(bool /*allowDataDependency*/) {
AIDGE_ASSERT(i == outputDimsOrder()[i],
"Permutation vector ({}) must be the identity above the input tensor rank ({}).", outputDimsOrder(), getInput(0)->dims());
}
+
mOutputs[0]->resize(outputDims);
+
+ // Set the data format of the output tensor
+ std::array<std::size_t, 5> outputDimsArray;
+ std::copy_n(outputDimsOrder().begin(), 5, outputDimsArray.begin());
+
+ mOutputs[0]->setDataFormat(getTransposeOutputDataFormat(getInput(0)->dataFormat(), outputDimsArray));
return true;
}
return false;
diff --git a/unit_tests/operator/Test_TransposeImpl.cpp b/unit_tests/operator/Test_TransposeImpl.cpp
index be54b8c426416c44a3587b2bb16c29b087e57e6d..0e3037773dae0a6f4c40e7ec1b4a45644ce8b017 100644
--- a/unit_tests/operator/Test_TransposeImpl.cpp
+++ b/unit_tests/operator/Test_TransposeImpl.cpp
@@ -205,7 +205,7 @@ TEST_CASE("[cpu/operator] Transpose DataFormat") {
SECTION("NCHW to NHWC"){
std::shared_ptr<Tensor> input = std::make_shared<Tensor>(std::vector<std::size_t>({16,3,224,450}));
input->setDataFormat(Aidge::DataFormat::NCHW);
-
+ input->setBackend("cpu");
auto transposeNode = Transpose({0,2,3,1});
auto op = std::static_pointer_cast<OperatorTensor>(transposeNode -> getOperator());
@@ -218,15 +218,99 @@ TEST_CASE("[cpu/operator] Transpose DataFormat") {
SECTION("NHWC to NCHW"){
std::shared_ptr<Tensor> input = std::make_shared<Tensor>(std::vector<std::size_t>({16,3,224,450}));
input->setDataFormat(Aidge::DataFormat::NHWC);
-
+ input->setBackend("cpu");
auto transposeNode = Transpose({0,3,1,2});
auto op = std::static_pointer_cast<OperatorTensor>(transposeNode -> getOperator());
op->associateInput(0,input);
op->setBackend("cpu");
transposeNode->forward();
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::NCHW);
+ }
+ SECTION("CHWN to NCHW") {
+ std::shared_ptr<Tensor> input = std::make_shared<Tensor>(std::vector<std::size_t>({3,224,450,16}));
+ input->setDataFormat(Aidge::DataFormat::CHWN);
+ input->setBackend("cpu");
+ auto transposeNode = Transpose({3,0,1,2});
- REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::NHWC);
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->setBackend("cpu");
+ transposeNode->forward();
+
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::NCHW);
+ }
+
+ SECTION("NCDHW to NDHWC") {
+ std::shared_ptr<Tensor> input = std::make_shared<Tensor>(std::vector<std::size_t>({16,3,8,224,450}));
+ input->setDataFormat(Aidge::DataFormat::NCDHW);
+ input->setBackend("cpu");
+ auto transposeNode = Transpose({0,2,3,4,1});
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->setBackend("cpu");
+ transposeNode->forward();
+
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::NDHWC);
+ }
+
+ SECTION("NDHWC to NCDHW") {
+ std::shared_ptr<Tensor> input = std::make_shared<Tensor>(std::vector<std::size_t>({16,8,224,450,3}));
+ input->setDataFormat(Aidge::DataFormat::NDHWC);
+ input->setBackend("cpu");
+ auto transposeNode = Transpose({0,4,1,2,3});
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->setBackend("cpu");
+ transposeNode->forward();
+
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::NCDHW);
+ }
+
+ SECTION("CDHWN to NCDHW") {
+ std::shared_ptr<Tensor> input = std::make_shared<Tensor>(std::vector<std::size_t>({3,8,224,450,16}));
+ input->setDataFormat(Aidge::DataFormat::CDHWN);
+ input->setBackend("cpu");
+ auto transposeNode = Transpose({4,0,1,2,3});
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->setBackend("cpu");
+ transposeNode->forward();
+
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::NCDHW);
}
+
+ SECTION("(Identity) NCHW to NCHW)") {
+ std::shared_ptr<Tensor> input = std::make_shared<Tensor>(std::vector<std::size_t>({16,3,224,450}));
+ input->setDataFormat(Aidge::DataFormat::NCHW);
+ input->setBackend("cpu");
+ auto transposeNode = Transpose({0,1,2,3}); // Identity permutation
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->setBackend("cpu");
+ transposeNode->forward();
+
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::NCHW);
+ }
+
+ SECTION("Invalid Format Handling") {
+ std::shared_ptr<Tensor> input = std::make_shared<Tensor>(std::vector<std::size_t>({16,3,224,450}));
+ input->setDataFormat(Aidge::DataFormat::Default);
+ input->setBackend("cpu");
+ auto transposeNode = Transpose({0,2,3,1});
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->setBackend("cpu");
+
+ transposeNode->forward();
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::Default);
+ //Should throw a Warning
+ }
+
}