Update on "[NT] Backward support for broadcasting binary ops"

soulitzer · soulitzer · commit 08d2e23154c0 · 2023-11-06T15:22:51.000-05:00
[ghstack-poisoned]
diff --git a/torch/csrc/autograd/input_metadata.cpp b/torch/csrc/autograd/input_metadata.cpp
@@ -9,8 +9,7 @@ namespace autograd {
 namespace {
 
 MetadataShape compute_variant_shape(const at::Tensor& input) {
-  if (input.is_nested() &&
-    !input.unsafeGetTensorImpl()->is_python_dispatch()) {
+  if (input.is_nested() && !input.unsafeGetTensorImpl()->is_python_dispatch()) {
     auto nested_size = input._nested_tensor_size();
     return MetadataShape{std::in_place_type<at::Tensor>, nested_size};
   }
@@ -80,7 +79,7 @@ at::Tensor InputMetadata::reduce_grad(at::Tensor& grad) const {
 std::stringstream InputMetadata::incompatible_shape_error_message(
     const size_t index,
     const at::Tensor& grad) const {
-  std::stringstream ss;
+  std::stringstream ss{};
   ss << "invalid gradient at index " << index << " - got ";
   if (::torch::autograd::is_cpp_nested_tensor(grad)) {
     ss << grad._nested_tensor_size();
@@ -97,6 +96,7 @@ std::stringstream InputMetadata::incompatible_shape_error_message(
 }
 
 bool InputMetadata::is_cpp_nested_tensor() const {
+  // NOLINTNEXTLINE(cppcoreguidelines-init-variables)
   bool ret = std::holds_alternative<at::Tensor>(shape_);
   TORCH_INTERNAL_ASSERT(ret == (is_nested_ && !is_tensor_subclass_))
   return ret;
@@ -114,20 +114,19 @@ SymIntSmallVec& InputMetadata::mutable_shape_as_dim_vector() {
 
 bool InputMetadata::is_nestedness_same(const at::Tensor& grad) const {
   return (
-    grad.is_nested() == is_nested_ &&
-    ::torch::autograd::is_cpp_nested_tensor(grad) == is_cpp_nested_tensor()
-  );
+      grad.is_nested() == is_nested_ &&
+      ::torch::autograd::is_cpp_nested_tensor(grad) == is_cpp_nested_tensor());
 }
 
 at::Tensor InputMetadata::shape_as_tensor() const {
   return std::get<at::Tensor>(shape_);
 }
 
 bool InputMetadata::maybe_expandable_to(const at::Tensor& grad) const {
-  // This is the initial step to determine whether or not the tensor represented by
-  // input_metadata is expandable to grad based on is-nestedness information alone.
-  // If this function returns true, then is_expandable_to_shape will be called.
-  // We support the following 3 types of expansion:
+  // This is the initial step to determine whether or not the tensor represented
+  // by input_metadata is expandable to grad based on is-nestedness information
+  // alone. If this function returns true, then is_expandable_to_shape will be
+  // called. We support the following 3 types of expansion:
   bool grad_is_nested = grad.is_nested();
   if (!is_nested_ && !grad_is_nested) {
     // Normal case (no NestedTensors are involved)
@@ -137,8 +136,8 @@ bool InputMetadata::maybe_expandable_to(const at::Tensor& grad) const {
     // (2) python NT -> python NT
     // (3) plain Tensor -> python NT
     return (
-      grad_is_nested && is_python_dispatch(grad) && (!is_nested_ || is_tensor_subclass_)
-    );
+        grad_is_nested && is_python_dispatch(grad) &&
+        (!is_nested_ || is_tensor_subclass_));
   }
 }
 
diff --git a/torch/csrc/autograd/input_metadata.h b/torch/csrc/autograd/input_metadata.h
@@ -38,7 +38,11 @@ using MetadataShape = std::variant<SymIntSmallVec, at::Tensor>;
  */
 struct TORCH_API InputMetadata {
   InputMetadata() = default;
-  InputMetadata(const at::TensorOptions& options, MetadataShape input_shape, bool is_tensor_subclass, bool is_nested);
+  InputMetadata(
+      const at::TensorOptions& options,
+      MetadataShape input_shape,
+      bool is_tensor_subclass,
+      bool is_nested);
   InputMetadata(const at::Tensor& t);
 
   const at::TensorOptions& options() const {
diff --git a/torch/nested/_internal/ops.py b/torch/nested/_internal/ops.py
@@ -46,11 +46,7 @@ def _wrap_jagged_dims(ndim, dims, op_name):
             " dimension is not supported for NestedTensor"
         )
     return (
-        tuple(
-            _outer_to_inner_dim(ndim, d)
-            for d in dims
-            if d != 0
-        ),
+        tuple(_outer_to_inner_dim(ndim, d) for d in dims if d != 0),
         zero_in_dims,
     )
 
