Function slice_tensor

fastdeploy/model_executor/models/tp_utils.py:222–230

slice_tensor

(tensor, start, end)

            return tensor.get_shape() if hasattr(tensor, "get_shape") else tensor.shape

        def slice_tensor(tensor, start, end):
            """slice_tensor"""
            shape = get_shape(tensor)
            if len(shape) == 1:
                return tensor[start:end]
            elif is_column:
                return tensor[..., start:end]
            else:
                return tensor[start:end, ...]

        q_end = num_attention_heads * head_dim
        k_end = q_end + num_key_value_heads * head_dim