diff --git a/layout/generic/ReflowInput.cpp b/layout/generic/ReflowInput.cpp
index e4b3732e7e68b..79fb5c094b619 100644
--- a/layout/generic/ReflowInput.cpp
+++ b/layout/generic/ReflowInput.cpp
@@ -2371,20 +2371,7 @@ void ReflowInput::InitConstraints(
           alignCB = alignCBParent;
         }
       }
-      if (alignCB->IsGridContainerFrame()) {
-        // Shrink-wrap grid items that will be aligned (rather than stretched)
-        // in its inline axis.
-        auto inlineAxisAlignment =
-            wm.IsOrthogonalTo(cbwm)
-                ? mStylePosition->UsedAlignSelf(alignCB->Style())._0
-                : mStylePosition->UsedJustifySelf(alignCB->Style())._0;
-        if ((inlineAxisAlignment != StyleAlignFlags::STRETCH &&
-             inlineAxisAlignment != StyleAlignFlags::NORMAL) ||
-            mStyleMargin->mMargin.GetIStart(wm).IsAuto() ||
-            mStyleMargin->mMargin.GetIEnd(wm).IsAuto()) {
-          mComputeSizeFlags += ComputeSizeFlag::ShrinkWrap;
-        }
-      } else {
+      if (!alignCB->IsGridContainerFrame()) {
         // Shrink-wrap blocks that are orthogonal to their container.
         if (isBlockLevel && mCBReflowInput &&
             mCBReflowInput->GetWritingMode().IsOrthogonalTo(mWritingMode)) {
diff --git a/layout/reftests/css-grid/grid-item-table-stretch-002-ref.html b/layout/reftests/css-grid/grid-item-table-stretch-002-ref.html
index 9bb60d5d84b2e..96909b1a035a9 100644
--- a/layout/reftests/css-grid/grid-item-table-stretch-002-ref.html
+++ b/layout/reftests/css-grid/grid-item-table-stretch-002-ref.html
@@ -32,7 +32,6 @@
   background: lightgrey;
 }
 caption { border: 1px dashed blue; }
-caption { justify-self: stretch; } /* XXX: Workaround bug 1350037. */
 x { display:block; width:16px; height:16px; }
 t { display:block; width:6px; height:6px; }
 
diff --git a/layout/tables/nsTableWrapperFrame.cpp b/layout/tables/nsTableWrapperFrame.cpp
index 87ac2e1899514..a3e957c4fd975 100644
--- a/layout/tables/nsTableWrapperFrame.cpp
+++ b/layout/tables/nsTableWrapperFrame.cpp
@@ -9,6 +9,7 @@
 #include "mozilla/ComputedStyle.h"
 #include "mozilla/PresShell.h"
 #include "nsFrameManager.h"
+#include "nsGridContainerFrame.h"
 #include "nsTableFrame.h"
 #include "nsTableCellFrame.h"
 #include "nsStyleConsts.h"
@@ -430,20 +431,7 @@ LogicalSize nsTableWrapperFrame::ComputeAutoSize(
   // actual size of the table, which (if it is specified as a percent)
   // could be something that is not reflected in our GetMinISize and
   // GetPrefISize.  See bug 349457 for an example.
-
-  // Shrink-wrap aChildFrame by default, except if we're a stretched grid item.
-  ComputeSizeFlags flags(ComputeSizeFlag::ShrinkWrap);
-  if (MOZ_UNLIKELY(IsGridItem()) && !StyleMargin()->HasInlineAxisAuto(aWM)) {
-    const auto* parent = GetParent();
-    auto inlineAxisAlignment =
-        aWM.IsOrthogonalTo(parent->GetWritingMode())
-            ? StylePosition()->UsedAlignSelf(parent->Style())._0
-            : StylePosition()->UsedJustifySelf(parent->Style())._0;
-    if (inlineAxisAlignment == StyleAlignFlags::NORMAL ||
-        inlineAxisAlignment == StyleAlignFlags::STRETCH) {
-      flags -= ComputeSizeFlag::ShrinkWrap;
-    }
-  }
+  const ComputeSizeFlags flags = CreateComputeSizeFlagsForChild();
 
   // Match the logic in Reflow() that sets aside space for the caption.
   Maybe<StyleCaptionSide> captionSide = GetCaptionSide();
@@ -555,6 +543,17 @@ void nsTableWrapperFrame::GetInnerOrigin(const MaybeCaptionSide& aCaptionSide,
   }
 }
 
+ComputeSizeFlags nsTableWrapperFrame::CreateComputeSizeFlagsForChild() const {
+  // Shrink-wrap child frames by default, except if we're a stretched grid item.
+  if (MOZ_UNLIKELY(IsGridItem())) {
+    auto* gridContainer = static_cast<nsGridContainerFrame*>(GetParent());
+    if (gridContainer->GridItemShouldStretch(this, eLogicalAxisInline)) {
+      return {};
+    }
+  }
+  return {ComputeSizeFlag::ShrinkWrap};
+}
+
 void nsTableWrapperFrame::CreateReflowInputForInnerTable(
     nsPresContext* aPresContext, nsTableFrame* aTableFrame,
     const ReflowInput& aOuterRI, Maybe<ReflowInput>& aChildRI,
@@ -588,11 +587,13 @@ void nsTableWrapperFrame::CreateReflowInputForInnerTable(
     }
   }
 
+  ComputeSizeFlags csFlags = CreateComputeSizeFlagsForChild();
   if (!aTableFrame->IsBorderCollapse() &&
       !aOuterRI.mStyleSizeOverrides.HasAnyOverrides()) {
     // We are not border-collapsed and not given any size overrides. It's
     // sufficient to call the standard ReflowInput constructor.
-    aChildRI.emplace(aPresContext, aOuterRI, aTableFrame, availSize, cbSize);
+    aChildRI.emplace(aPresContext, aOuterRI, aTableFrame, availSize, cbSize,
+                     ReflowInput::InitFlags{}, StyleSizeOverrides{}, csFlags);
     return;
   }
 
@@ -618,7 +619,8 @@ void nsTableWrapperFrame::CreateReflowInputForInnerTable(
       aBSizeOccupiedByCaption);
 
   aChildRI.emplace(aPresContext, aOuterRI, aTableFrame, availSize, Nothing(),
-                   ReflowInput::InitFlag::CallerWillInit, innerOverrides);
+                   ReflowInput::InitFlag::CallerWillInit, innerOverrides,
+                   csFlags);
   aChildRI->Init(aPresContext, cbSize, Some(*borderPadding - *padding),
                  padding);
 }
diff --git a/layout/tables/nsTableWrapperFrame.h b/layout/tables/nsTableWrapperFrame.h
index 5696f254400af..4e4c03e60792e 100644
--- a/layout/tables/nsTableWrapperFrame.h
+++ b/layout/tables/nsTableWrapperFrame.h
@@ -5,6 +5,7 @@
 #ifndef nsTableWrapperFrame_h__
 #define nsTableWrapperFrame_h__
 
+#include "LayoutConstants.h"
 #include "mozilla/Attributes.h"
 #include "mozilla/Maybe.h"
 #include "nscore.h"
@@ -203,6 +204,15 @@ class nsTableWrapperFrame : public nsContainerFrame {
                       mozilla::LogicalPoint& aOrigin,
                       mozilla::WritingMode aWM) const;
 
+  // This is a helper for CreateReflowInputForInnerTable() and
+  // ComputeAutoSize(). It computes whether we need shrink-wrap behavior for
+  // children.
+  //
+  // Note: We don't need to call this in CreateReflowInputForCaption() because
+  // when we reflow the captions, we want them to stretch their inline-sizes to
+  // be at least as wide as the inner table frame.
+  mozilla::ComputeSizeFlags CreateComputeSizeFlagsForChild() const;
+
   // Create and init the child reflow input, using passed-in aChildRI, so that
   // caller can use it after we return.
   //