Fix std and var for float16 and float32 (#186)

* Fix std and var for float16 and float32
- fix type mismatch
- add correction parameter to bessel correction calculation
- use unbiased (correction=1) std / var by default

* remove space

Co-authored-by: Ronian <ronian@Ronians-MBP.attlocal.net>

Author: 2 people

aten/src/ATen/native/mps/operations/ReduceOps.mm

@@ -744,8 +744,8 @@ Tensor std_var_common_impl_mps(
     }
   }
 
-  bool use_correction = correction.has_value();
-  const auto correction_value = use_correction ? correction.value() : false;
+  bool use_correction = !(correction.has_value() && correction.value() == 0);
+  const auto correction_value = correction.value_or(1);
   int64_t correction_n = 1;
 
   native_mps::MPSGraphCache* cache_ = native_mps::MPSGraphCache::getInstance();
@@ -884,7 +884,7 @@ Tensor std_var_common_impl_mps(
      return output_t;
   }
 
-  double bessel_correction = ((double) correction_n) / ((double) (correction_n-1));
+  double bessel_correction = ((double) correction_n) / ((double) (correction_n-correction_value));
 
   auto stream = at::mps::getCurrentMPSStream();
 
@@ -894,7 +894,7 @@ Tensor std_var_common_impl_mps(
     string bessel_corrected = (use_correction && correction_value) ? "unbiased " : "biased ";
     string use_dim_info = (use_dim) ? "use_dim=1:" + to_string(dim_value.size()) : "use_dim=0";
     string keepdim_info = (keepdim) ? "keepdim=1" : "keepdim=0";
-    string key = op_key + use_dim_info + ":" + keepdim_info + ":" + string([ns_key UTF8String]) + ":" + native_mps::getTensorsStringKey(input_t) + ":" + bessel_corrected;
+    string key = op_key + use_dim_info + ":" + keepdim_info + ":" + string([ns_key UTF8String]) + ":" + native_mps::getTensorsStringKey(input_t) + ":" + bessel_corrected + ":" + std::to_string(correction_value);
 
     auto cachedGraph = cache_->LookUpAs<CachedGraph>(key);
     // Initialize once if configuration not found in cache
@@ -916,7 +916,7 @@ Tensor std_var_common_impl_mps(
           if (use_correction && correction_value)
           {
               MPSGraphTensor *besselTensor= [mpsGraph constantWithScalar:bessel_correction
-                                                    dataType:MPSDataTypeFloat32];
+                                                    dataType: native_mps::getMPSDataType(input_t.scalar_type())];
               MPSGraphTensor *correctedTensor = [mpsGraph multiplicationWithPrimaryTensor: outputVarTensor
                                                                           secondaryTensor: besselTensor
                                                                                      name: nil];

test/test_mps.py

@@ -8300,6 +8300,8 @@ class TestConsistency(TestCase):
         'linalg.cross': ['f16', 'f32', 'i16', 'i32', 'i64', 'u8'],
         'unique_consecutive': ['b8', 'f16', 'f32', 'i16', 'i32', 'i64', 'u8'],
         'nn.functional.nll_loss': ['f32'],
+        'std': ['f16','f32'],
+        'var': ['f16','f32'],
     }
 
 
@@ -8485,9 +8487,7 @@ class TestConsistency(TestCase):
         'masked.sum': [torch.bool],
 
         # Functions that hard crash
-        'std': [torch.float16],
         'stft': [torch.float32],
-        'var': [torch.float16],
         # + forward when requires_grad=True or running backward
         '__rpow__': [torch.int64],
         'masked.std': [torch.int32],