chore: bubble attribute error up

exla/c_src/exla/exla_mlir.cc

@@ -121,14 +121,7 @@ mlir::Type MLIRModule::ParseType(std::string string) {
 }
 
 mlir::Attribute MLIRModule::ParseAttribute(std::string string) {
-  auto attribute = mlir::parseAttribute(string, context_);
-
-  if (attribute == nullptr) {
-    std::cerr << "Unable to parse MLIR attribute: " << string << std::endl;
-    exit(1);
-  }
-
-  return attribute;
+  return mlir::parseAttribute(string, context_);
 }
 
 }  // namespace exla

exla/c_src/exla/exla_mlir.h

@@ -58,7 +58,7 @@ class MLIRModule {
 
   std::string ToString();
 
-  // Note: returns nullptr if the parsing fails
+  // Note: ParseAttribute and ParseType return nullptr if the parsing fails
   mlir::Type ParseType(std::string);
   mlir::Attribute ParseAttribute(std::string);
 

exla/lib/exla/mlir/value.ex

@@ -750,11 +750,21 @@ defmodule EXLA.MLIR.Value do
 
     call_target_name =
       case op_type do
-        {:f, 32} -> "qr_cpu_custom_call_f32"
-        {:f, 64} -> "qr_cpu_custom_call_f64"
-        {:f, 16} -> "qr_cpu_custom_call_f16"
-        {:bf, 16} -> "qr_cpu_custom_call_bf16"
-        type -> raise "QR decomposition not supported for type #{inspect(type)}"
+        {:f, 32} ->
+          "qr_cpu_custom_call_f32"
+
+        {:f, 64} ->
+          "qr_cpu_custom_call_f64"
+
+        {:f, 16} ->
+          "qr_cpu_custom_call_f16"
+
+        {:bf, 16} ->
+          "qr_cpu_custom_call_bf16"
+
+        type ->
+          # Due to matching on EXLA.Defn, we are sure that the device here is always :host
+          raise "QR decomposition not supported on :host device for type #{inspect(type)}"
       end
 
     attributes = [