diff --git a/src/com/android/launcher3/widget/WidgetsBottomSheet.java b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
index c347939..ceb0072 100644
--- a/src/com/android/launcher3/widget/WidgetsBottomSheet.java
+++ b/src/com/android/launcher3/widget/WidgetsBottomSheet.java
@@ -16,6 +16,7 @@
 
 package com.android.launcher3.widget;
 
+import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_BOTTOM_WIDGETS_TRAY;
 
 import android.content.Context;
@@ -187,7 +188,13 @@
                 mWidgetCellHorizontalPadding)
                 .forEach(row -> {
                     TableRow tableRow = new TableRow(getContext());
-                    tableRow.setGravity(Gravity.TOP);
+                    if (enableCategorizedWidgetSuggestions()) {
+                        // Vertically center align items, so that even if they don't fill bounds,
+                        // they can look organized when placed together in a row.
+                        tableRow.setGravity(Gravity.CENTER_VERTICAL);
+                    } else {
+                        tableRow.setGravity(Gravity.TOP);
+                    }
                     row.forEach(widgetItem -> {
                         WidgetCell widget = addItemCell(tableRow);
                         widget.applyFromCellItem(widgetItem);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
index 4e704fd..2e51fb6 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsFullSheet.java
@@ -162,7 +162,7 @@
     @Nullable PersonalWorkPagedView mViewPager;
     private boolean mIsInSearchMode;
     private boolean mIsNoWidgetsViewNeeded;
-    @Px private int mMaxSpanPerRow;
+    @Px protected int mMaxSpanPerRow;
     protected DeviceProfile mDeviceProfile;
 
     protected TextView mNoWidgetsView;
diff --git a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
index c7d2aa3..f10ab48 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsListTableViewHolderBinder.java
@@ -15,6 +15,8 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
+
 import android.content.Context;
 import android.graphics.Bitmap;
 import android.util.Log;
@@ -147,7 +149,13 @@
                 tableRow = (TableRow) table.getChildAt(i);
             } else {
                 tableRow = new TableRow(table.getContext());
-                tableRow.setGravity(Gravity.TOP);
+                if (enableCategorizedWidgetSuggestions()) {
+                    // Vertically center align items, so that even if they don't fill bounds, they
+                    // can look organized when placed together in a row.
+                    tableRow.setGravity(Gravity.CENTER_VERTICAL);
+                } else {
+                    tableRow.setGravity(Gravity.TOP);
+                }
                 table.addView(tableRow);
             }
             if (tableRow.getChildCount() > widgetItems.size()) {
diff --git a/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java b/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
index 06cc65e..2d17033 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsRecommendationTableLayout.java
@@ -15,6 +15,7 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
 import static com.android.launcher3.LauncherSettings.Favorites.CONTAINER_WIDGETS_PREDICTION;
 
 import android.content.Context;
@@ -104,8 +105,13 @@
         for (int i = 0; i < data.mRecommendationTable.size(); i++) {
             List<WidgetItem> widgetItems = data.mRecommendationTable.get(i);
             TableRow tableRow = new TableRow(getContext());
-            tableRow.setGravity(Gravity.TOP);
-
+            if (enableCategorizedWidgetSuggestions()) {
+                // Vertically center align items, so that even if they don't fill bounds, they can
+                // look organized when placed together in a row.
+                tableRow.setGravity(Gravity.CENTER_VERTICAL);
+            } else {
+                tableRow.setGravity(Gravity.TOP);
+            }
             for (WidgetItem widgetItem : widgetItems) {
                 WidgetCell widgetCell = addItemCell(tableRow);
                 widgetCell.applyFromCellItem(widgetItem, data.mPreviewScale);
diff --git a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
index 744c45b..4326515 100644
--- a/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
+++ b/src/com/android/launcher3/widget/picker/WidgetsTwoPaneSheet.java
@@ -15,6 +15,7 @@
  */
 package com.android.launcher3.widget.picker;
 
+import static com.android.launcher3.Flags.enableCategorizedWidgetSuggestions;
 import static com.android.launcher3.Flags.enableUnfoldedTwoPanePicker;
 
 import android.content.Context;
@@ -308,15 +309,25 @@
                 if (mSuggestedWidgetsHeader != null) {
                     mSuggestedWidgetsHeader.setExpanded(false);
                 }
+
+                WidgetsListContentEntry contentEntryToBind;
+                if (enableCategorizedWidgetSuggestions()) {
+                    // Setting max span size enables row to understand how to fit more than one item
+                    // in a row.
+                    contentEntryToBind = contentEntry.withMaxSpanSize(mMaxSpanPerRow);
+                } else {
+                    contentEntryToBind = contentEntry;
+                }
+
                 WidgetsRowViewHolder widgetsRowViewHolder =
                         mWidgetsListTableViewHolderBinder.newViewHolder(mRightPane);
                 mWidgetsListTableViewHolderBinder.bindViewHolder(widgetsRowViewHolder,
-                        contentEntry,
+                        contentEntryToBind,
                         ViewHolderBinder.POSITION_FIRST | ViewHolderBinder.POSITION_LAST,
                         Collections.EMPTY_LIST);
                 widgetsRowViewHolder.mDataCallback = data -> {
                     mWidgetsListTableViewHolderBinder.bindViewHolder(widgetsRowViewHolder,
-                            contentEntry,
+                            contentEntryToBind,
                             ViewHolderBinder.POSITION_FIRST | ViewHolderBinder.POSITION_LAST,
                             Collections.singletonList(data));
                 };