Add dtype functions for ops that take dtype from 1st operand

build_tools/python_deploy/build_linux_packages.sh

@@ -288,7 +288,7 @@ function test_in_tree() {
   python -m e2e_testing.main --config=lazy_tensor_core -v
 
   echo ":::: Run TorchDynamo e2e integration tests"
-  python -m e2e_testing.main --config=torchdynamo -v --crashing_tests_to_not_attempt_to_run_and_a_bug_is_filed RandnDtypeDeviceModule_basic RandnLikeModule_basic Matmul_dot
+  python -m e2e_testing.main --config=torchdynamo -v --crashing_tests_to_not_attempt_to_run_and_a_bug_is_filed RandnDtypeDeviceModule_basic RandnLikeModule_basic
 }
 
 function setup_venv() {

e2e_testing/xfail_sets.py

@@ -96,8 +96,6 @@
 
     # Dtype function transition failures
     "MobilenetV3Module_basic",
-    "ResNet18Module_basic",
-    "ResNet18StaticModule_basic",
 }
 
 STABLEHLO_PASS_SET = {
@@ -706,7 +704,8 @@
     "FullLikeModuleInt2DStatic_basic",
     "FullModuleInt3D_basic",
     "FullModuleFloat2D_basic",
-    "RepeatModule_basic"
+    "RepeatModule_basic",
+    "ResNet18StaticModule_basic",
 }
 
 LTC_XFAIL_SET = {

lib/Dialect/Torch/Transforms/RefineTypes.cpp

@@ -613,39 +613,6 @@ void TypeAnalysis::visitOperation(Operation *op,
     return;
   }
 
-  // Take dtype from first operand.
-  if (isa<CopyToValueTensorOp, CopyToNonValueTensorOp, AtenBatchNormOp,
-          AtenReluOp, AtenRelu6Op, AtenGeluOp, AtenCeilOp, AtenGeluBackwardOp,
-          AtenBitwiseNotOp, AtenToPrimDeviceOp, AtenCpuOp, AtenContiguousOp,
-          AtenDetachOp, AtenMaskedFill_ScalarOp, AtenCopyOp, AtenCumsumOp,
-          AtenLayerNormOp, AtenClampOp, AtenClampMinOp, AtenClampMaxOp,
-          AtenNegOp, AtenFloorOp, Aten_SoftmaxBackwardDataOp, AtenDropoutOp,
-          AtenTanhBackwardOp, AtenHardtanhBackwardOp,
-          Aten_LogSoftmaxBackwardDataOp, AtenAddIntOp, AtenAbsOp,
-          AtenThresholdOp, AtenSquareOp, AtenUniformOp, AtenBernoulliOp,
-          AtenBernoulli_FloatOp, AtenBernoulliTensorOp,
-          ValsemVariantAtenBernoulliFloatOp, AtenBernoulliTensorOp,
-          AtenBernoulliPOp, AtenFillScalarOp, AtenHardsigmoidOp, AtenCloneOp,
-          AtenHardswishOp, AtenSiluOp, AtenHardtanhOp, AtenMaskedSelectOp,
-          AtenMaxPool2dOp, AtenAvgPool2dOp, AtenAdaptiveAvgPool2dOp,
-          AtenFlattenUsingIntsOp, AtenSqueezeOp, AtenSqueezeDimOp,
-          AtenUnsqueezeOp, AtenViewOp, Aten_UnsafeViewOp, AtenReshapeOp,
-          Aten_ReshapeAliasOp, AtenResize_Op, AtenTransposeIntOp, AtenTOp,
-          AtenPermuteOp, AtenIndexSelectOp, AtenSelectIntOp,
-          AtenSelectScatterOp, AtenNarrowOp, AtenSliceTensorOp,
-          AtenScatterReduceTwoOp, AtenSliceScatterOp, AtenGatherOp,
-          AtenExpandOp, AtenExpandAsOp, AtenBroadcastToOp, AtenRepeatOp,
-          AtenConstantPadNdOp, AtenPadOp, AtenZero_Op, AtenIndexTensorOp,
-          Aten_IndexPutImplOp, AtenIndexPutOp, AtenCopyOp, AtenZeroOp,
-          AtenIndexPutHackedTwinOp, AtenPreluOp, AtenMaskedFillScalarOp,
-          AtenFlipOp, PrimAbsScalarOp, AtenNumpyTOp, AtenTriuOp,
-          AtenMaskedFillTensorOp, AtenRollOp, AtenPowTensorTensorOp,
-          AtenLiftFreshCopyOp, AtenIndexTensorHackedTwinOp,
-          AtenUpsampleNearest2dOp, AtenMishOp, AtenRoundOp, AtenFillTensorOp,
-          AtenUpsampleNearest2dBackwardOp, AtenLeakyReluBackwardOp>(op)) {
-    return incorporateKnowledge(op->getResult(0), operands[0]->getValue());
-  }
-
   // Dtype is always float32, except for bfloat16, float64 and nullptr after
   // promotion and assuming possible-zero rank.
   if (isa<AtenAtan2Op>(op)) {
@@ -1365,8 +1332,7 @@ static Type getMostRefinedStaticType(Value v, DataFlowSolver &solver) {
 // the right thing forthose ops.
 //
 static bool allowsTypeRefinementOrIsSafeToRefine(Operation *op) {
-  return op->hasTrait<mlir::torch::Torch::OpTrait::AllowsTypeRefinement>() ||
-         isa<CopyToNonValueTensorOp, CopyToValueTensorOp>(op);
+  return op->hasTrait<mlir::torch::Torch::OpTrait::AllowsTypeRefinement>();
 }
 
 // Some operations have extra verification logic regarding the relationship

lib/Dialect/Torch/Transforms/ReifyAbstractInterpCalculationsUtils.cpp

@@ -167,12 +167,14 @@ FailureOr<Value> Torch::adjustFunctionArg(
     return b.create<DerefineOp>(loc, desiredType, operand).getResult();
   }
 
-  // !torch.union<int, float> is the type used for `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.
+  // !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.
   if (auto unionType = desiredType.dyn_cast<Torch::UnionType>()) {
     if (llvm::all_of(unionType.getContainedTypes(), [](Type containedType) {
-          return containedType.isa<Torch::IntType, Torch::FloatType>();
+          return containedType
+              .isa<Torch::IntType, Torch::FloatType, Torch::NoneType>();
         }))
       return b.create<DerefineOp>(loc, desiredType, operand).getResult();
   }