Update list for supported MTIA sharding types (#2789)

Summary:
Pull Request resolved: #2789

As title.

Previous diffs in stack add unit tests for TWCW, TWRW, and GRID sharding.

bypass-github-pytorch-ci-checks

Reviewed By: kausv, nautsimon, egienvalue, rexyl

Differential Revision: D69924446

fbshipit-source-id: cf9da78f1492e22d0f01ab327f4aeeeb5398cc86

Author: jvandebon

torchrec/distributed/embedding_types.py

@@ -414,9 +414,16 @@ def __init__(
         self._fused_params = fused_params
 
     def sharding_types(self, compute_device_type: str) -> List[str]:
-        # For MTIA, sharding types are restricted to TW, CW.
+
         if compute_device_type in {"mtia"}:
-            return [ShardingType.TABLE_WISE.value, ShardingType.COLUMN_WISE.value]
+            return [
+                ShardingType.TABLE_WISE.value,
+                ShardingType.COLUMN_WISE.value,
+                ShardingType.TABLE_COLUMN_WISE.value,
+                ShardingType.ROW_WISE.value,
+                ShardingType.TABLE_ROW_WISE.value,
+                ShardingType.GRID_SHARD.value,
+            ]
 
         types = [
             ShardingType.DATA_PARALLEL.value,