Add an operation to compute the mean along a dimension. Computes the mean along the dimension 'dim' of the input tensor. It is implemented using the IReduceLayer from TensorRT. Parameters: input : Tensor The input tensor. dim : int The dim
(input: Tensor,
dim: Union[int, Tuple[int]],
keepdim: bool = False)
| 3199 | |
| 3200 | |
| 3201 | def mean(input: Tensor, |
| 3202 | dim: Union[int, Tuple[int]], |
| 3203 | keepdim: bool = False) -> Tensor: |
| 3204 | ''' |
| 3205 | Add an operation to compute the mean along a dimension. |
| 3206 | |
| 3207 | Computes the mean along the dimension 'dim' of the input tensor. |
| 3208 | |
| 3209 | It is implemented using the IReduceLayer from TensorRT. |
| 3210 | |
| 3211 | Parameters: |
| 3212 | input : Tensor |
| 3213 | The input tensor. |
| 3214 | |
| 3215 | dim : int |
| 3216 | The dimension along which the mean is computed. |
| 3217 | |
| 3218 | keepdim : bool |
| 3219 | Is the dimension kept in the reduced tensor? When True the |
| 3220 | dimension is kept, it is removed from the shape otherwise. |
| 3221 | |
| 3222 | Returns: |
| 3223 | The tensor produced by this reduction operation. |
| 3224 | ''' |
| 3225 | return reduce(input, op=trt.ReduceOperation.AVG, dim=dim, keepdim=keepdim) |
| 3226 | |
| 3227 | |
| 3228 | def max(input: Tensor, dim: int, keepdim: bool = False) -> Tensor: |