how to generate KeyedJaggedTensor on GPU

[xiexbing]

Hello,

I compressed the values, and lengths to reduce the PCIe cost. then I want to move the keys, compressed values and lengths to GPU, and then decompress and generate KJT on GPUs.

My question is: any examples on how to do this?

thanks.

[henrylhtsang]

what do you mean by "compress"? They should be torch.tensors right?

[xiexbing]

Yes, My question is: can I generate KJT on GPUs directly? The reason I want to do so is I want to reduce the data size transferred via PCIe, I need to convert the compressed tensors to the values/lengths tensors as KJT format.

…

On Tue, Feb 20, 2024 at 11:03 AM henrylhtsang ***@***.***> wrote: what do you mean by "compress"? They should be torch.tensors right? — Reply to this email directly, view it on GitHub <#1715 (comment)> or unsubscribe <https://github.com/notifications/unsubscribe-auth/AQGDI2HSVJE6I2CLGAW6JE3YUTXR3BFKMF2HI4TJMJ2XIZLTSOBKK5TBNR2WLJDUOJ2WLJDOMFWWLO3UNBZGKYLEL5YGC4TUNFRWS4DBNZ2F6YLDORUXM2LUPGBKK5TBNR2WLJDUOJ2WLJDOMFWWLLTXMF2GG2C7MFRXI2LWNF2HTAVFOZQWY5LFUVUXG43VMWSG4YLNMWVXI2DSMVQWIX3UPFYGLLDTOVRGUZLDORPXI6LQMWWES43TOVSUG33NNVSW45FGORXXA2LDOOJIFJDUPFYGLKTSMVYG643JORXXE6NFOZQWY5LFVEZTQNJUGA4DMNRTQKSHI6LQMWSWS43TOVS2K5TBNR2WLKRSGE2DCNRVHAZTEM5HORZGSZ3HMVZKMY3SMVQXIZI> . You are receiving this email because you authored the thread. Triage notifications on the go with GitHub Mobile for iOS <https://apps.apple.com/app/apple-store/id1477376905?ct=notification-email&mt=8&pt=524675> or Android <https://play.google.com/store/apps/details?id=com.github.android&referrer=utm_campaign%3Dnotification-email%26utm_medium%3Demail%26utm_source%3Dgithub> .

[xiexbing]

sorry, I have keys as a string list, values as a list tensor, lengths as a list tensor.

[henrylhtsang]

@xiexbing I think you can just first move the tensors to GPU, and then initialize a kjt that way? if I remember correctly, that should work just fine?