debug on KJT issue

Summary:
# context
* In the IR export workflow, the module takes KJTs as input and produces an `ExportedProgram` of the module
* KJT actually has a variable length for the values and weights
* This dynamic nature of KJT needs to be explicitly passed to torch.export

# changes
* add a util function to mark the input KJT's dynamic shape
* add in the test of how to correctly specify the dynamics shapes for the input KJT

# results
* input KJTs with different value lengths
```
(Pdb) feature1.values()
tensor([0, 1, 2, 3, 2, 3])
(Pdb) feature2.values()
tensor([0, 1, 2, 3, 2, 3, 4])
```
* exported_program can take those input KJTs
```
(Pdb) ep.module()(feature1)
[tensor([[-2.8735e-16, -2.8735e-16, -2.8735e-16, -2.8735e-16],
        [-2.8735e-16, -2.8735e-16, -2.8735e-16, -2.8735e-16]]), tensor([[-2.8735e-16, -2.8735e-16, -2.8735e-16, -2.8735e-16, -1.4368e-15,
         -1.4368e-15, -1.4368e-15, -1.4368e-15],
        [-2.8735e-16, -2.8735e-16, -2.8735e-16, -2.8735e-16, -1.4368e-15,
         -1.4368e-15, -1.4368e-15, -1.4368e-15]])]
(Pdb) ep.module()(feature2)
[tensor([[-2.8735e-16, -2.8735e-16, -2.8735e-16, -2.8735e-16],
        [-2.8735e-16, -2.8735e-16, -2.8735e-16, -2.8735e-16]]), tensor([[-2.8735e-16, -2.8735e-16, -2.8735e-16, -2.8735e-16, -1.4368e-15,
         -1.4368e-15, -1.4368e-15, -1.4368e-15],
        [-2.8735e-16, -2.8735e-16, -2.8735e-16, -2.8735e-16, -1.4368e-15,
         -1.4368e-15, -1.4368e-15, -1.4368e-15]])]
```
* deserialized module can take those input KJTs
```
(Pdb) deserialized_model(feature1)
[tensor([[ 0.2630,  0.1473, -0.3691,  0.2261],
        [ 0.0000,  0.0000,  0.0000,  0.0000]],
       grad_fn=<SplitWithSizesBackward0>), tensor([[ 0.2198, -0.1648, -0.0121,  0.1998, -0.0384, -0.2458, -0.6844,  0.8741],
        [ 0.1313,  0.2968, -0.2979, -0.2150, -0.2593,  0.6758,  1.0010,  0.9052]],
       grad_fn=<SplitWithSizesBackward0>)]
(Pdb) deserialized_model(feature2)
[tensor([[ 0.2630,  0.1473, -0.3691,  0.2261],
        [ 0.0000,  0.0000,  0.0000,  0.0000]],
       grad_fn=<SplitWithSizesBackward0>), tensor([[ 0.2198, -0.1648, -0.0121,  0.1998, -0.0384, -0.2458, -0.6844,  0.8741],
        [ 0.1313,  0.2968, -0.2979, -0.2150, -0.9359,  0.1123,  0.5834, -0.1357]],
       grad_fn=<SplitWithSizesBackward0>)]
```

Differential Revision: D57824907

Author: TroyGarden

torchrec/ir/utils.py

@@ -9,17 +9,21 @@
 
 #!/usr/bin/env python3
 
-from typing import List, Optional, Tuple, Type
+from collections import defaultdict
+from typing import Dict, List, Optional, Tuple, Type, Union
 
 import torch
 
 from torch import nn
-from torch.export.exported_program import ExportedProgram
+from torch.export import Dim, ExportedProgram, ShapesCollection
+from torch.export.dynamic_shapes import _Dim as DIM
+from torchrec import KeyedJaggedTensor
 from torchrec.ir.types import SerializerInterface
 
 
 # TODO: Replace the default interface with the python dataclass interface
 DEFAULT_SERIALIZER_CLS = SerializerInterface
+DYNAMIC_DIMS: Dict[str, int] = defaultdict(int)
 
 
 def serialize_embedding_modules(
@@ -88,3 +92,45 @@ def deserialize_embedding_modules(
         setattr(parent, attrs[-1], new_module)
 
     return model
+
+
+def _get_dim(x: Union[DIM, str, None], s: str) -> DIM:
+    if isinstance(x, DIM):
+        return x
+    elif isinstance(x, str):
+        if x in DYNAMIC_DIMS:
+            DYNAMIC_DIMS[x] += 1
+            x += str(DYNAMIC_DIMS[x])
+        dim = Dim(x)
+    else:
+        DYNAMIC_DIMS[s] += 1
+        dim = Dim(s + str(DYNAMIC_DIMS[s]))
+    return dim
+
+
+def mark_dynamic_kjt(
+    kjt: KeyedJaggedTensor,
+    vlen: Union[DIM, str, None] = None,
+    llen: Union[DIM, str, None] = None,
+    lofs: Union[DIM, str, None] = None,
+    variable_length: bool = False,
+    shapes: Optional[ShapesCollection] = None,
+) -> ShapesCollection:
+    """
+    Makes the given KJT dynamic.
+    """
+    global DYNAMIC_DIMS
+    if shapes is None:
+        shapes = ShapesCollection()
+    vlen = _get_dim(vlen, "vlen")
+    shapes[kjt._values] = (vlen,)
+    if kjt._weights is not None:
+        shapes[kjt._weights] = (vlen,)
+    if variable_length:
+        llen = _get_dim(llen, "llen")
+        lofs = _get_dim(lofs, "lofs")
+        if kjt._lengths is not None:
+            shapes[kjt._lengths] = (llen,)
+        if kjt._offsets is not None:
+            shapes[kjt._offsets] = (lofs,)
+    return shapes