Add torch check for dtype within bilinear (pytorch#118900)

Fixes pytorch#117237
Short-term fix, when dtype does not match, it will be reflected in the torch check.

@ezyang a cpp test case is added
Pull Request resolved: pytorch#118900
Approved by: https://github.com/ezyang, https://github.com/malfet

Author: lancerts

aten/src/ATen/native/Linear.cpp

@@ -703,6 +703,28 @@ Tensor bilinear(const Tensor& input1, const Tensor& input2, const Tensor& weight
   // See [Note: hacky wrapper removal for optional tensor]
   c10::MaybeOwned<Tensor> bias_maybe_owned = at::borrow_from_optional_tensor(bias_opt);
   const Tensor& bias = *bias_maybe_owned;
+  if (bias.defined()) {
+    TORCH_CHECK(
+        input1.dtype() == input2.dtype() && input1.dtype() == weight.dtype() &&
+            input1.dtype() == bias.dtype(),
+        "All tensors must have the same dtype, got input1: ",
+        input1.dtype(),
+        ", input2: ",
+        input2.dtype(),
+        ", weight: ",
+        weight.dtype(),
+        ", bias: ",
+        bias.dtype());
+  } else {
+    TORCH_CHECK(
+        input1.dtype() == input2.dtype() && input1.dtype() == weight.dtype(),
+        "All tensors must have the same dtype, got input1: ",
+        input1.dtype(),
+        ", input2: ",
+        input2.dtype(),
+        ", weight: ",
+        weight.dtype());
+  }
 
   TORCH_CHECK(input1.dim() == input2.dim(), "bilinear(): input dimensions do not match: got ", input1.dim(), " and ", input2.dim());
   for (const auto i : c10::irange(input1.dim() - 1)) {

test/cpp/api/functional.cpp

@@ -1524,6 +1524,13 @@ TEST_F(FunctionalTest, Bilinear) {
   ASSERT_EQ(y_no_bias.sizes(), torch::IntArrayRef({2, 1}));
   auto y_no_bias_exp = torch::tensor({{448, 1701}}).reshape({2, 1});
   ASSERT_TRUE(torch::allclose(y_no_bias, y_no_bias_exp, 1e-4, 1e-7));
+
+  input1 = input1.to(torch::kFloat64);
+  input2 = input2.to(torch::kInt32);
+  weight = weight.to(torch::kInt32);
+  ASSERT_THROWS_WITH(
+      F::bilinear(input1, input2, weight),
+      "All tensors must have the same dtype, got input1: double, input2: int, weight: int");
 }
 
 TEST_F(FunctionalTest, Normalize) {