Add an elementwise operation with two inputs. For each input, that function first creates a constant tensor if the input is an integer or a float. Then, if needed, it expands the smaller tensor to make sure its rank is the same as the larger one. Then, it performs the elementwi
(left: Union[Tensor, int,
float], right: Union[Tensor, int, float],
op: trt.ElementWiseOperation)
| 2959 | |
| 2960 | |
| 2961 | def elementwise_binary(left: Union[Tensor, int, |
| 2962 | float], right: Union[Tensor, int, float], |
| 2963 | op: trt.ElementWiseOperation) -> Tensor: |
| 2964 | ''' |
| 2965 | Add an elementwise operation with two inputs. |
| 2966 | |
| 2967 | For each input, that function first creates a constant tensor if the input |
| 2968 | is an integer or a float. Then, if needed, it expands the smaller tensor to |
| 2969 | make sure its rank is the same as the larger one. Then, it performs the |
| 2970 | elementwise operation 'op'. |
| 2971 | |
| 2972 | The following closures are defined in functional.*: |
| 2973 | |
| 2974 | add for op=trt.ElementWiseOperation.SUM |
| 2975 | sub for op=trt.ElementWiseOperation.SUB |
| 2976 | mul for op=trt.ElementWiseOperation.PROD |
| 2977 | div for op=trt.ElementWiseOperation.DIV |
| 2978 | floordiv for op=trt.ElementWiseOperation.FLOOR_DIV |
| 2979 | gt for op=trt.ElementWiseOperation.GREATER |
| 2980 | lt for op=trt.ElementWiseOperation.LESS |
| 2981 | op_and for op=trt.ElementWiseOperation.AND |
| 2982 | op_or for op=trt.ElementWiseOperation.OR |
| 2983 | eq for op=trt.ElementWiseOperation.EQUAL |
| 2984 | minimum for op=trt.ElementWiseOperation.MIN |
| 2985 | maximum for op=trt.ElementWiseOperation.MAX |
| 2986 | pow for op=trt.ElementWiseOperation.POW |
| 2987 | |
| 2988 | It is implemented using the IElementWiseLayer from TensorRT. |
| 2989 | |
| 2990 | Parameters: |
| 2991 | left : Union[Tensor, int, float] |
| 2992 | The first input. If that input is an integer or a float, the |
| 2993 | function creates a constant tensor. |
| 2994 | |
| 2995 | right : Union[Tensor, int, float] |
| 2996 | The second input. If that input is an integer or a float, the |
| 2997 | function creates a constant tensor. |
| 2998 | |
| 2999 | op : trt.ElementWiseOperation |
| 3000 | The binary operation to perform. |
| 3001 | |
| 3002 | Returns: |
| 3003 | The tensor produced by this elementwise operation. |
| 3004 | ''' |
| 3005 | left, right = broadcast_helper(left, right) |
| 3006 | if left.dtype == trt.int32 and right.dtype == trt.int64: |
| 3007 | left = cast(left, trt.int64) |
| 3008 | if left.dtype == trt.int64 and right.dtype == trt.int32: |
| 3009 | right = cast(right, trt.int64) |
| 3010 | layer = default_trtnet().add_elementwise(left.trt_tensor, right.trt_tensor, |
| 3011 | op) |
| 3012 | return _create_tensor(layer.get_output(0), layer) |
| 3013 | |
| 3014 | |
| 3015 | add = partial(elementwise_binary, op=trt.ElementWiseOperation.SUM) |
nothing calls this directly
no test coverage detected