Slicing tensor with dynamic shape in custom enc-dec architecture

[AvivSham]

Hi,
As mentioned in this issue #2535 we are looking to convert a custom version of enc-dec architecture. We want to pass the encoder output to the decoder in some way, to do so we need to slice some tensors in the dec. However, these tensors have at least one dynamic dimension (which is the batch size dimension).
What is the correct way to slice a tensor with a dynamic shape?

@yuekaizhang maybe you can help here?

Cheers,

[yuekaizhang]

Sorry, I don't know.

Maybe you could get some insights from how we implement the code here https://github.com/k2-fsa/icefall/blob/master/egs/aishell/ASR/whisper/whisper_encoder_forward_monkey_patch.py#L15.

TRT_LLM: https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/models/enc_dec/model.py#L1977, which used a position_ids inputs to avoid the slice on the tensors with dynamic shapes.

[AvivSham]

@yuekaizhang Thank you for your quick response.
Unfortunately, both links are irrelevant for our case as we need to slice output tensors of intermediate decoder layers. Could you please direct us to someone who can assist us?

[AvivSham]

I have a follow-up question let's say that we fix all dimensions including the batch in prepare_inputs function. what is the correct way to slice a tensor?

1. use slice from tensorrt_llm.functional and use it as part of the module's forward function.
2. Add the relevant slice layer to the network using add_slice method.