Added support for slogdet in LazyTensor shape inference (#77904)

Fixes #3576

Added support for `slogdet` in LazyTensor shape inference
Pull Request resolved: pytorch/pytorch#77904
Approved by: https://github.com/wconstab, https://github.com/JackCaoG

torch/csrc/lazy/core/shape_inference.cpp

@@ -443,6 +443,17 @@ std::vector<Shape> compute_shape_smooth_l1_loss(
   }
 }
 
+std::vector<Shape> compute_shape_slogdet(const at::Tensor & self) {
+  // assumes self.shape is {*, n, n} and returns shape *
+  TORCH_INTERNAL_ASSERT(self.dim() >= 2);
+  std::vector<int64_t> out_sizes(self.sizes().begin(), self.sizes().end() - 2);
+  // Doesn't check input dtype, but output dtype either matches it,
+  // or the actual slogdet operation will throw if it's an unsupported type.
+  // Sign and det outputs hold the same shape, dtype.
+  return {Shape(self.scalar_type(), out_sizes),
+          Shape(self.scalar_type(), out_sizes)};
+}
+
 std::vector<Shape> compute_shape_smooth_l1_loss_backward(
     const at::Tensor& grad_output, const at::Tensor& self,
     const at::Tensor& target, int64_t reduction, double beta) {

torch/csrc/lazy/core/shape_inference.h

@@ -56,6 +56,7 @@ TORCH_API std::vector<torch::lazy::Shape> compute_shape_random_functional(const
 TORCH_API std::vector<torch::lazy::Shape> compute_shape_random_functional(const at::Tensor & self, int64_t from, c10::optional<int64_t> to, c10::optional<at::Generator> generator);
 TORCH_API std::vector<torch::lazy::Shape> compute_shape_relu(const at::Tensor & self);
 TORCH_API std::vector<torch::lazy::Shape> compute_shape_repeat(const at::Tensor & self, at::IntArrayRef repeats);
+TORCH_API std::vector<torch::lazy::Shape> compute_shape_slogdet(const at::Tensor & self);
 TORCH_API std::vector<torch::lazy::Shape> compute_shape_smooth_l1_loss_backward(const at::Tensor & grad_output, const at::Tensor & self, const at::Tensor & target, int64_t reduction, double beta);
 TORCH_API std::vector<torch::lazy::Shape> compute_shape_sort(const at::Tensor & self, int64_t dim, bool descending);
 TORCH_API std::vector<torch::lazy::Shape> compute_shape_stack(at::TensorList tensors, int64_t dim);