Merge branch 'main' of https://github.com/CPestka/daphne

CPestka · Dec 14, 2023 · 0ff5b6f · 0ff5b6f
2 parents 296bb14 + 91bfa36
commit 0ff5b6f
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/src/ir/daphneir/DaphneInferTypesOpInterface.cpp b/src/ir/daphneir/DaphneInferTypesOpInterface.cpp
@@ -197,6 +197,22 @@ std::vector<Type> daphne::CreateFrameOp::inferTypes() {
 }
 
 std::vector<Type> daphne::RandMatrixOp::inferTypes() {
+    auto elTy = getMin().getType();
+    if (elTy == UnknownType::get(getContext())) {
+        elTy = getMax().getType();
+    } else {
+        assert((getMax().getType() == UnknownType::get(getContext()) || elTy == getMax().getType()) &&
+               "Min and max need to have the same type");
+    }
+    return {daphne::MatrixType::get(getContext(), elTy)};
+}
+
+std::vector<Type> daphne::EigenOp::inferTypes() {
+    auto evMatType = getArg().getType().dyn_cast<daphne::MatrixType>();
+    return {evMatType.withSameElementType(), evMatType};
+}
+
+std::vector<Type> daphne::RandTensor3DOp::inferTypes() {
     auto elTy = getMin().getType();
     if(elTy == UnknownType::get(getContext())) {
         elTy = getMax().getType();