Register serialized name to KJT/JT (#1713)

Summary:

Trying to reland D51312977. Hopefully this wont break any torch
packages as D53139358 recently did a similar change where called
`register_pytree_node` with an extra new argument, `flatten_with_keys_fn`,
which was added recently in D52547850.

Differential Revision: D53857843

torchrec/sparse/jagged_tensor.py

@@ -7,6 +7,7 @@
 
 import abc
 import operator
+import inspect
 from typing import Any, Callable, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -635,7 +636,11 @@ def _jt_flatten_spec(t: JaggedTensor, spec: TreeSpec) -> List[Optional[torch.Ten
 
 
 register_pytree_node(
-    JaggedTensor, _jt_flatten, _jt_unflatten, flatten_with_keys_fn=_jt_flatten_with_keys
+    JaggedTensor,
+    _jt_flatten,
+    _jt_unflatten,
+    flatten_with_keys_fn=_jt_flatten_with_keys,
+    serialized_type_name="torchrec.JaggedTensor",
 )
 register_pytree_flatten_spec(JaggedTensor, _jt_flatten_spec)
 
@@ -2094,6 +2099,7 @@ def _kjt_flatten_spec(
     _kjt_flatten,
     _kjt_unflatten,
     flatten_with_keys_fn=_kjt_flatten_with_keys,
+    serialized_type_name="torchrec.KeyedJaggedTensor",
 )
 register_pytree_flatten_spec(KeyedJaggedTensor, _kjt_flatten_spec)