**select_input** This API takes in multiple inputs and uses an integer mask to select one input to output. It is useful in control flow. Args: inputs(tuple|list): The input variables mask(Tensor): A tensor containing 1 integer number selecting which inp
(inputs, mask)
| 2026 | |
| 2027 | |
| 2028 | def select_input(inputs, mask): |
| 2029 | """ |
| 2030 | **select_input** |
| 2031 | |
| 2032 | This API takes in multiple inputs and uses an integer mask to select one |
| 2033 | input to output. It is useful in control flow. |
| 2034 | |
| 2035 | Args: |
| 2036 | inputs(tuple|list): The input variables |
| 2037 | mask(Tensor): A tensor containing 1 integer number selecting which |
| 2038 | input to output |
| 2039 | |
| 2040 | Returns: |
| 2041 | Variable: The selected input variable |
| 2042 | """ |
| 2043 | helper = LayerHelper('select_input', **locals()) |
| 2044 | check_type(inputs, 'inputs', (list, tuple), 'select_input') |
| 2045 | check_variable_and_dtype(mask, 'mask', ['int32'], 'select_input') |
| 2046 | |
| 2047 | # Select input should expand the shape. If it is - 1 and valid number, use - 1 first. If the dim is different, an error will be reported directly |
| 2048 | # assert inputs[0].dtype == inputs[1].dtype, f"Expect the inputs should have the same dtype, but get {inputs[0].dtype} and {inputs[1].dtype}" |
| 2049 | |
| 2050 | output_shape = _select_input_infer_shape(inputs[0].shape, inputs[1].shape) |
| 2051 | output_dtype = inputs[1].dtype |
| 2052 | output_type = inputs[1].type |
| 2053 | |
| 2054 | out = helper.create_variable( |
| 2055 | dtype=output_dtype, shape=output_shape, type=output_type |
| 2056 | ) |
| 2057 | helper.append_op( |
| 2058 | type='select_input', |
| 2059 | inputs={'X': inputs, 'Mask': mask}, |
| 2060 | outputs={'Out': out}, |
| 2061 | ) |
| 2062 | return out |
| 2063 | |
| 2064 | |
| 2065 | def select_input_with_buildin_type(inputs, mask, name): |
no test coverage detected