diff --git a/include/aidge/data/DataFormat.hpp b/include/aidge/data/DataFormat.hpp
index 4df854a1ea880e3e7def7ff6aee798ab95d30254..d89eb24cdc7d37c1bd58a1e1e97024a9f5bd47b4 100644
--- a/include/aidge/data/DataFormat.hpp
+++ b/include/aidge/data/DataFormat.hpp
@@ -83,6 +83,24 @@ 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 input format.
+ * 2. Applies the given output dimensions order to compute the new permutation.
+ * 3. Identifies the corresponding data format after applying 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.
+ *
+ * @pre `outputDimsOrder` must be a valid permutation of the format's dimensions.
+ */
+DataFormat getTransposedDataFormat(const DataFormat inputDataFormat, const std::array<std::size_t, 5> outputDimsOrder);
+
} // namespace Aidge
namespace {
diff --git a/python_binding/data/pybind_DataFormat.cpp b/python_binding/data/pybind_DataFormat.cpp
index 5308fb3023e15a74f0dc5f674917d7ae65cbb52f..354911c33ba8357b2cdbc18c62237b60f1ec5d74 100644
--- a/python_binding/data/pybind_DataFormat.cpp
+++ b/python_binding/data/pybind_DataFormat.cpp
@@ -67,6 +67,7 @@ void init_DataFormat(py::module& m) {
bindEnum<DataFormat>(m, "dformat");
m.def("format_as", (const char* (*)(DataFormat)) &format_as, py::arg("df"));
m.def("get_permutation_mapping", &getPermutationMapping, py::arg("src"), py::arg("dst"));
+ m.def("get_transposed_data_format", &getTransposedDataFormat, py::arg("input_data_format"), py::arg("output_dims_order"));
}
} // namespace Aidge
diff --git a/src/data/DataFormat.cpp b/src/data/DataFormat.cpp
index 8b7460b3d42de2e2ddb60a6332405141f87b297c..095ef9368e13b1d31322e73d98d1ec63ab0fda3b 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 ``DataFormatTranspose``.
+ *
+ * @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 getTransposedDataFormat(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..f9e46947654a1117891dbe1e5cd679ae0e436dec 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
+ DataFormatTranspose outputDimsArray;
+ std::copy_n(outputDimsOrder().begin(), outputDimsOrder().size(), outputDimsArray.begin());
+
+ mOutputs[0]->setDataFormat(getTransposedDataFormat(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 02b338dd265ceb4d9e03bbf86398cc8db38045df..f852c3fe0ebf2e6c73dddd358bf778b67779bffc 100644
--- a/unit_tests/operator/Test_TransposeImpl.cpp
+++ b/unit_tests/operator/Test_TransposeImpl.cpp
@@ -200,3 +200,113 @@ TEST_CASE("[cpu/operator] Transpose(forward)") {
REQUIRE(*(op->getOutput(0)) == *output);
}
}
+
+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);
+ auto transposeNode = Transpose({0,2,3,1});
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode -> getOperator());
+ op->associateInput(0,input);
+ op->forwardDims();
+
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::NHWC);
+ }
+ 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);
+ auto transposeNode = Transpose({0,3,1,2});
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode -> getOperator());
+ op->associateInput(0,input);
+ op->forwardDims();
+ 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);
+ auto transposeNode = Transpose({3,0,1,2});
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->forwardDims();
+
+ 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);
+ auto transposeNode = Transpose({0,2,3,4,1});
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->forwardDims();
+
+ 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);
+ auto transposeNode = Transpose({0,4,1,2,3});
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->forwardDims();
+
+ 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);
+ auto transposeNode = Transpose({4,0,1,2,3});
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->forwardDims();
+
+
+ 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);
+ auto transposeNode = Transpose({0,1,2,3}); // Identity permutation
+
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+ op->forwardDims();
+
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::NCHW);
+ }
+
+ SECTION("Default to Default") {
+ std::shared_ptr<Tensor> input = std::make_shared<Tensor>(std::vector<std::size_t>({16,3,224,450}));
+ input->setDataFormat(Aidge::DataFormat::Default);
+ auto transposeNode = Transpose({0,2,3,1});
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+
+ op->forwardDims();
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::Default);
+ //Should print a Log::Warning
+ }
+
+ 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::NCHW);
+ auto transposeNode = Transpose({0,2,1,3});
+ auto op = std::static_pointer_cast<OperatorTensor>(transposeNode->getOperator());
+ op->associateInput(0, input);
+
+ op->forwardDims();
+ REQUIRE(op->getOutput(0)->dataFormat() == DataFormat::Default);
+ //Should print a Log::Warning
+ }
+
+}
+