Fix handling of !torch.number in abstract interpretation library (#…

…2309)

In PyTorch, the `NumberType` is equal to `Union[int, float,
complex]`. However, the abstract interpretation library was treating
the `NumberType` as `Union[int, float]`, resulting in type mismatches
when reifying certain dtype functions. This commit fixes the type
inconsistency by having the abstract interpretation functions take as
an input a `Union[int, float, complex]` for the ops that take
`!torch.number` inputs.

lib/Dialect/Torch/Transforms/ReifyAbstractInterpCalculationsUtils.cpp

@@ -176,10 +176,17 @@ FailureOr<Value> Torch::adjustFunctionArg(
     return b.create<DerefineOp>(loc, desiredType, operand).getResult();
   }
 
-  // !torch.union<int, float> or !torch.union<int, float, none> is the type used
-  // for (optional) `Scalar` inputs. At compile time, such inputs will usually
-  // be resolved to an `int` or a `float` so we need to derefine to match the
-  // library function signature.
+  // The type `!torch.number` can be an `int`, `float`, or `complex`.
+  // TODO: Add a new type `Torch::ComplexType` to handle the complex case.
+  if (desiredType.isa<Torch::NumberType>() &&
+      operandType.isa<Torch::IntType, Torch::FloatType>()) {
+    return b.create<DerefineOp>(loc, desiredType, operand).getResult();
+  }
+
+  // !torch.union<int, float, none> is the type used for optional
+  // `Scalar` inputs. At compile time, such inputs will usually be
+  // resolved to an `int`, `float`, or `None` so we need to derefine
+  // to match the library function signature.
   if (auto unionType = desiredType.dyn_cast<Torch::UnionType>()) {
     if (llvm::all_of(unionType.getContainedTypes(), [](Type containedType) {
           return containedType