[Documentation] more accurate discription of behavior of KeyedJaggedTensor

[gouchangjiang]

I found an interesting phenomenon that could be enhanced when using KeyedJaggedTensor.

import torch
from torchrec.sparse.jagged_tensor import JaggedTensor, KeyedJaggedTensor

values = [
    torch.Tensor([1.0]),
    torch.Tensor(),
    torch.Tensor([7.0, 8.0]),
    torch.Tensor([10.0, 11.0, 12.0]),
]
weights = [
    torch.Tensor([1.0]),
    torch.Tensor(),
    torch.Tensor([7.0, 8.0]),
    torch.Tensor([10.0, 11.0, 12.0]),
]
j1 = JaggedTensor.from_dense(
    values=values,
    weights=weights,
)

kjt1 = KeyedJaggedTensor.from_jt_dict({'j1': j1})

# first print
for k in kjt1.keys():
  print(f'{k}, {kjt1[k]}')

# second print
for v in kjt1.values():
  print(v)

# third print
print(kjt1['j1'])

# 4-th print
print(kjt1['j1'].values())

The second print does not meet my expectation, and does not align well with the documentation.

Returns the values of the KeyedJaggedTensor.

The third printing exactly follows the documentation. Maybe we can say that the method values() of KJT lists all values of, with each element as a torch.tensor. Or I misunderstood the concept, the values of the KJT, is not a JaggedTensor, but flattened values of all JaggedTensors.