Update on "Fix kthvalue error for scalar input"

fixes #30818

Note that the median case was already fixed by #45847

[ghstack-poisoned]

.circleci/config.yml

@@ -655,11 +655,13 @@ jobs:
           echo "Retrieving test reports"
           docker cp $id:/var/lib/jenkins/workspace/test/test-reports ./ || echo 'No test reports found!'
           if [[ ${BUILD_ENVIRONMENT} == *"coverage"* ]]; then
+              echo "Retrieving C++ coverage report"
+              docker cp $id:/var/lib/jenkins/workspace/build/coverage.info ./test
+          fi
+          if [[ ${BUILD_ENVIRONMENT} == *"coverage"* || ${BUILD_ENVIRONMENT} == *"onnx"* ]]; then
               echo "Retrieving Python coverage report"
               docker cp $id:/var/lib/jenkins/workspace/test/.coverage ./test
               docker cp $id:/var/lib/jenkins/workspace/test/coverage.xml ./test
-              echo "Retrieving C++ coverage report"
-              docker cp $id:/var/lib/jenkins/workspace/build/coverage.info ./test
               python3 -mpip install codecov
               python3 -mcodecov
           fi

.circleci/verbatim-sources/job-specs/pytorch-job-specs.yml

@@ -217,11 +217,13 @@ jobs:
           echo "Retrieving test reports"
           docker cp $id:/var/lib/jenkins/workspace/test/test-reports ./ || echo 'No test reports found!'
           if [[ ${BUILD_ENVIRONMENT} == *"coverage"* ]]; then
+              echo "Retrieving C++ coverage report"
+              docker cp $id:/var/lib/jenkins/workspace/build/coverage.info ./test
+          fi
+          if [[ ${BUILD_ENVIRONMENT} == *"coverage"* || ${BUILD_ENVIRONMENT} == *"onnx"* ]]; then
               echo "Retrieving Python coverage report"
               docker cp $id:/var/lib/jenkins/workspace/test/.coverage ./test
               docker cp $id:/var/lib/jenkins/workspace/test/coverage.xml ./test
-              echo "Retrieving C++ coverage report"
-              docker cp $id:/var/lib/jenkins/workspace/build/coverage.info ./test
               python3 -mpip install codecov
               python3 -mcodecov
           fi

.jenkins/caffe2/test.sh

@@ -93,6 +93,7 @@ fi
   # See comments on
   # https://github.com/HypothesisWorks/hypothesis-python/commit/eadd62e467d6cee6216e71b391951ec25b4f5830
   $MAYBE_SUDO pip -q uninstall -y hypothesis
+  $MAYBE_SUDO pip -q uninstall -y coverage
   # "pip install hypothesis==3.44.6" from official server is unreliable on
   # CircleCI, so we host a copy on S3 instead
   $MAYBE_SUDO pip -q install attrs==18.1.0 -f https://s3.amazonaws.com/ossci-linux/wheels/attrs-18.1.0-py2.py3-none-any.whl

.jenkins/pytorch/codegen-test.sh

@@ -42,6 +42,7 @@ python -m tools.pyi.gen_pyi \
 # autograd codegen (called by torch codegen but can run independently)
 python -m tools.autograd.gen_autograd \
   "$OUT"/torch/share/ATen/Declarations.yaml \
+  aten/src/ATen/native/native_functions.yaml \
   "$OUT"/autograd \
   tools/autograd
 

aten/src/ATen/core/boxing/impl/make_boxed_from_unboxed_functor.h

@@ -323,8 +323,6 @@ namespace impl {
   call_functor_with_args_from_stack_(Functor* functor, Stack* stack, std::index_sequence<ivalue_arg_indices...>) {
     (void)(stack); // when sizeof...(ivalue_arg_indices) == 0, this argument would be unused and we have to silence the compiler warning.
 
-    constexpr size_t num_ivalue_args = sizeof...(ivalue_arg_indices);
-
     /*
      * For ops that take "Tensor&" as an argument, ivalue_to_arg would still return a "Tensor" by value
      * and C++ doesn't allow us to call (*functor) with a temporary "Tensor" when it expects "Tensor&".
@@ -335,7 +333,7 @@ namespace impl {
     using ArgTypes = typename guts::infer_function_traits_t<Functor>::parameter_types;
     return (*functor)(reference_cast<guts::typelist::element_t<ivalue_arg_indices, ArgTypes>>(
       ivalue_to_arg<std::decay_t<guts::typelist::element_t<ivalue_arg_indices, ArgTypes>>, AllowDeprecatedTypes>::call(
-        std::move(torch::jit::peek(*stack, ivalue_arg_indices, num_ivalue_args))
+        std::move(torch::jit::peek(*stack, ivalue_arg_indices, sizeof...(ivalue_arg_indices)))
     ))...);
   }
 

aten/src/ATen/native/ForeachOpsKernels.cpp

@@ -149,20 +149,48 @@ void foreach_tensor_##OP##_scalarlist_slow_(TensorList input, TensorList tensors
 
 FOREACH_BINARY_OP_LIST_ALPHA(add);
 FOREACH_BINARY_OP_LIST_ALPHA(sub);
+
 FOREACH_BINARY_OP_SCALAR(add);
 FOREACH_BINARY_OP_SCALAR(sub);
 FOREACH_BINARY_OP_SCALAR(mul);
 FOREACH_BINARY_OP_SCALAR(div);
+
 FOREACH_BINARY_OP_SCALARLIST(add);
 FOREACH_BINARY_OP_SCALARLIST(sub);
 FOREACH_BINARY_OP_SCALARLIST(mul);
 FOREACH_BINARY_OP_SCALARLIST(div);
+
 FOREACH_BINARY_OP_LIST(mul);
 FOREACH_BINARY_OP_LIST(div);
+
 FOREACH_UNARY_OP(sqrt);
 FOREACH_UNARY_OP(exp);
+FOREACH_UNARY_OP(abs);
+FOREACH_UNARY_OP(acos);
+FOREACH_UNARY_OP(asin);
+FOREACH_UNARY_OP(atan);
+FOREACH_UNARY_OP(ceil);
+FOREACH_UNARY_OP(cos);
+FOREACH_UNARY_OP(cosh);
+FOREACH_UNARY_OP(erf);
+FOREACH_UNARY_OP(erfc);
+FOREACH_UNARY_OP(expm1);
+FOREACH_UNARY_OP(floor);
+FOREACH_UNARY_OP(log);
+FOREACH_UNARY_OP(log10);
+FOREACH_UNARY_OP(log1p);
+FOREACH_UNARY_OP(log2);
+FOREACH_UNARY_OP(neg);
+FOREACH_UNARY_OP(tan);
+FOREACH_UNARY_OP(tanh);
+FOREACH_UNARY_OP(sin);
+FOREACH_UNARY_OP(sinh);
+FOREACH_UNARY_OP(round);
+FOREACH_UNARY_OP(lgamma);
+
 FOREACH_POINTWISE_OP_SCALAR(addcdiv);
 FOREACH_POINTWISE_OP_SCALAR(addcmul);
+
 FOREACH_POINTWISE_OP_SCALARLIST(addcdiv);
 FOREACH_POINTWISE_OP_SCALARLIST(addcmul);
 

aten/src/ATen/native/Resize.h

@@ -1,6 +1,7 @@
 #pragma once
 
 #include <ATen/ATen.h>
+#include <ATen/native/ResizeCommon.h>
 #include <TH/THTensor.hpp>
 
 namespace at { namespace native {
@@ -51,15 +52,7 @@ inline TensorImpl* resize_impl_cpu_(
   if (stride) {
     self->set_sizes_and_strides(size, *stride);
     // NB: storage size can be different from numel.
-    for (size_t dim = 0; dim < size.size(); ++dim) {
-      // FIXME: Don't rely on storage_size being negative because this
-      // may not be true for some edge cases.
-      if (size[dim] == 0) {
-        storage_size = 0;
-        break;
-      }
-      storage_size += (size[dim] - 1) * stride.value()[dim];
-    }
+    storage_size = storage_size_for(size, *stride);
   } else {
     self->set_sizes_contiguous(size);
     storage_size = self->numel();

aten/src/ATen/native/ResizeCommon.h

@@ -5,6 +5,21 @@
 
 namespace at { namespace native {
 
+inline int64_t storage_size_for(IntArrayRef size, IntArrayRef stride) {
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(size.size() == stride.size(),
+      "storage_size_for(size, stride) requires that size and stride ",
+      "have the same size as a precondition.");
+  int64_t storage_size = 1;
+  for (size_t dim = 0; dim < size.size(); ++dim) {
+    if (size[dim] == 0) {
+      storage_size = 0;
+      break;
+    }
+    storage_size += (size[dim] - 1) * stride[dim];
+  }
+  return storage_size;
+}
+
 inline Tensor& resize_named_tensor_(
     Tensor& self,
     IntArrayRef size,

aten/src/ATen/native/TensorFactories.cpp

@@ -341,6 +341,15 @@ Tensor new_empty(
   return at::empty(size, self.options().merge_in(options));
 }
 
+Tensor new_empty_strided(
+    const Tensor& self,
+    IntArrayRef size,
+    IntArrayRef stride,
+    const TensorOptions& options
+    ) {
+  return at::empty_strided(size, stride, self.options().merge_in(options));
+}
+
 // ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ eye ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 Tensor eye(int64_t n, const TensorOptions& options) {

aten/src/ATen/native/TensorShape.cpp

@@ -1868,7 +1868,7 @@ Tensor diag(const Tensor& self, int64_t dimension) {
 }
 
 Tensor& diag_cpu_out(Tensor &result, const Tensor& self, int64_t dimension) {
-  AT_DISPATCH_ALL_TYPES_AND(at::ScalarType::Bool, self.scalar_type(), "diag", [&] {
+  AT_DISPATCH_ALL_TYPES_AND_COMPLEX_AND(at::ScalarType::Bool, self.scalar_type(), "diag", [&] {
     apply_diag<scalar_t>(result, self, dimension);
   });
   return result;