Function conv1d

tensorrt_llm/functional.py:3546–3583

(input: Tensor,
           weight: Tensor,
           bias: Optional[Tensor] = None,
           stride: int = 1,
           padding: int = 0,
           dilation: int = 1,
           groups: int = 1)

def conv1d(input: Tensor,
           weight: Tensor,
           bias: Optional[Tensor] = None,
           stride: int = 1,
           padding: int = 0,
           dilation: int = 1,
           groups: int = 1) -> Tensor:

    noutput = weight.size()[0]
    kernel_size = weight.size()[-2]
    is_weight_constant = (weight.producer is not None
                          and weight.producer.type == trt.LayerType.CONSTANT)
    weight = weight.producer.weights if is_weight_constant else trt.Weights()

    if bias is not None:
        is_bias_constant = (bias.producer is not None
                            and bias.producer.type == trt.LayerType.CONSTANT)
        bias = bias.producer.weights if is_bias_constant else trt.Weights()

    input_shuffled = stack([input], dim=input.ndim())
    kernel_size = trt.Dims([kernel_size, 1])

    layer = default_trtnet().add_convolution_nd(input_shuffled.trt_tensor,
                                                noutput, kernel_size, weight,
                                                bias)
    layer.stride_nd = (stride, 2)
    layer.padding_nd = (padding, 0)
    layer.dilation_nd = (dilation, 2)
    layer.num_groups = groups

    if not is_weight_constant:
        layer.set_input(1, weight.trt_tensor)
    if bias is not None and not is_bias_constant:
        layer.set_input(2, bias.trt_tensor)

    output_2d = _create_tensor(layer.get_output(0), layer)
    output_1d = squeeze(output_2d, dim=-1)
    return output_1d


def conv2d(input: Tensor,