Fix unary min / max and permutation shapes

test/cpp/cpp_test_util.cpp

@@ -26,6 +26,12 @@ at::Tensor ToCpuTensor(const at::Tensor& t) {
 }
 
 bool EqualValues(at::Tensor tensor1, at::Tensor tensor2) {
+  if (tensor1.sizes() != tensor2.sizes()) {
+    std::cerr << "Different sizes:\n"
+              << tensor1.sizes() << "\n-vs-\n"
+              << tensor2.sizes() << "\n";
+    return false;
+  }
   tensor1 = ToCpuTensor(tensor1);
   tensor2 = ToCpuTensor(tensor2);
 
@@ -45,11 +51,17 @@ void ForEachDevice(const std::function<void(const Device&)>& devfn) {
 
 bool CloseValues(at::Tensor tensor1, at::Tensor tensor2, double rtol,
                  double atol) {
+  if (tensor1.sizes() != tensor2.sizes()) {
+    std::cerr << "Different sizes:\n"
+              << tensor1.sizes() << "\n-vs-\n"
+              << tensor2.sizes() << "\n";
+    return false;
+  }
   tensor1 = ToCpuTensor(tensor1);
   tensor2 = ToCpuTensor(tensor2);
   bool equal = tensor1.allclose(tensor2, rtol, atol);
   if (!equal) {
-    std::cout << tensor1 << "\n-vs-\n" << tensor2 << "\n";
+    std::cerr << tensor1 << "\n-vs-\n" << tensor2 << "\n";
   }
   return equal;
 }

torch_xla/csrc/helpers.h

@@ -168,6 +168,23 @@ class XlaHelpers {
                                   xla::int64 expected_rank,
                                   xla::int64 offset = 0);
 
+  // Gathers the input using the order specified by the permutation. For each i,
+  // output[i] = input[permutation[i]]. The given permutation must be the same
+  // size as the input.
+  template <typename Container>
+  static std::vector<typename Container::value_type> Permute(
+      tensorflow::gtl::ArraySlice<const xla::int64> permutation,
+      const Container& input) {
+    using T = typename Container::value_type;
+    XLA_CHECK(xla::IsPermutation(permutation, input.size()))
+        << "Invalid permutation specified";
+    std::vector<T> output(input.size());
+    for (size_t i = 0; i < permutation.size(); ++i) {
+      output[i] = input[permutation[i]];
+    }
+    return output;
+  }
+
   // Creates a transposition from the given input and dimensions.
   static std::vector<xla::int64> MakeTransposePermutation(xla::int64 dim0,
                                                           xla::int64 dim1,

torch_xla/csrc/ops/ops.cpp

@@ -559,7 +559,7 @@ NodePtr MaxUnary(const Value& input) {
     return node.ReturnOp(xla::Reshape(result, {1}), loctx);
   };
   return GenericOp(OpKind(at::aten::max), {input},
-                   xla::ShapeUtil::MakeShape(input.shape().element_type(), {1}),
+                   xla::ShapeUtil::MakeShape(input.shape().element_type(), {}),
                    std::move(lower_fn));
 }
 
@@ -577,7 +577,7 @@ NodePtr MinUnary(const Value& input) {
     return node.ReturnOp(xla::Reshape(result, {1}), loctx);
   };
   return GenericOp(OpKind(at::aten::min), {input},
-                   xla::ShapeUtil::MakeShape(input.shape().element_type(), {1}),
+                   xla::ShapeUtil::MakeShape(input.shape().element_type(), {}),
                    std::move(lower_fn));
 }
 

torch_xla/csrc/view.cpp

@@ -8,6 +8,7 @@
 #include "tensorflow/compiler/xla/util.h"
 #include "tensorflow/compiler/xla/xla_client/debug_macros.h"
 #include "tensorflow/compiler/xla/xla_client/util.h"
+#include "torch_xla/csrc/helpers.h"
 #include "torch_xla/csrc/ops/generic_slice.h"
 #include "torch_xla/csrc/ops/permute.h"
 #include "torch_xla/csrc/ops/select.h"
@@ -78,7 +79,8 @@ ViewInfo::ViewInfo(xla::Shape shape, std::vector<xla::int64> sizes)
 
 ViewInfo::ViewInfo(std::vector<xla::int64> sizes,
                    std::vector<xla::int64> permutation, xla::PrimitiveType type)
-    : shape(xla::ShapeUtil::MakeShape(type, xla::Permute(permutation, sizes))),
+    : shape(xla::ShapeUtil::MakeShape(type,
+                                      XlaHelpers::Permute(permutation, sizes))),
       sizes(std::move(sizes)),
       permutation(std::move(permutation)) {}