JaggedTensor (TorchRec)
This API is an open-source API of the TorchRec and is not an external API of the Rec SDK Torch. This section describes the parameter ranges supported by the TorchRec APIs called when the Rec SDK Torch is used.
Function
Class that holds sparse IDs and feature lengths, which is used to query tables. For example, values are [id1, id2, id3, id4] and lengths are [1, 2, 1]. Indicates that the embeddings of id2 and id3 after table lookup should be pooled.
Prototype
1 2 | class JaggedTensor: def __init__(**kwargs): |
Parameters
Parameter |
Type |
Mandatory/Optional |
Description |
|---|---|---|---|
values |
torch.Tensor[int64] |
Mandatory |
ID of the sparse table. |
weights |
torch.Tensor |
Optional |
When an NPU is used, the value can only be the default value None and cannot be customized. |
lengths |
torch.Tensor[int64] |
Optional |
Length of the feature sequence in each sample. This parameter is mandatory when an NPU is used. The value range is [1, 10000]. Ensure that the sum of lengths is equal to the length of values. NOTICE:
Currently, the Rec SDK Torch does not support variable batch size. The lengths of all JaggedTensors in a training job must be the same. |
offsets |
torch.Tensor[int64] |
Optional |
The offsets value is the accumulated result of lengths. The first digit of offsets is 0, and the subsequent digits are the accumulated values of lengths. The default value is None.
NOTICE:
The validity of offsets is ensured by the user. |
Example
1 2 | from torchrec import JaggedTensor JaggedTensor(values=[1, 3, 4], lengths=[1, 1, 1], offsets=[0, 1, 2, 3]) |
See Also
For details about the API call sequence and example, see Porting and Training.