Coverage for mlprodict/onnx_conv/sklconv/function_transformer_converters.py: 96%

1""" 

2@file 

3@brief Rewrites some of the converters implemented in 

4:epkg:`sklearn-onnx`. 

5""" 

6import copy 

7from onnx.helper import make_tensor 

8from skl2onnx.common.data_types import guess_numpy_type 

9from skl2onnx.common._apply_operation import apply_concat, apply_identity 

10from ...onnx_tools.onnx2py_helper import ( 

11 _var_as_dict, guess_proto_dtype, get_tensor_shape) 

12from ...npy.onnx_version import FctVersion 

13 

14 

15def new_calculate_sklearn_function_transformer_output_shapes(operator): 

16 """ 

17 Rewrites the converters implemented in 

18 :epkg:`sklearn-onnx` to support custom functions 

19 implemented with :ref:`l-numpy-onnxpy`. 

20 """ 

21 fct = operator.raw_operator.func 

22 if hasattr(fct, 'signed_compiled'): 

23 dtype = guess_numpy_type(operator.inputs[0].type) 

24 fct = fct[FctVersion((dtype, ), None)] 

25 if hasattr(fct, 'compiled'): 

26 compiled = fct.compiled 

27 if not hasattr(compiled, 'onnx_'): 

28 raise RuntimeError( # pragma: no cover 

29 "Attribute 'onnx_' is missing, function was not " 

30 "converted to onnx.") 

31 onx = compiled.onnx_ 

32 graph = onx.graph 

33 outputs = graph.output 

34 

35 # Let's assume there is only one output 

36 # with the same type as the input. 

37 # Only the shape changes. 

38 if len(outputs) != 1: 

39 raise RuntimeError( # pragma: no cover 

40 f"Only one output is allowed not {len(outputs)}.") 

41 input_type = operator.inputs[0].type.__class__ 

42 if compiled.meta_.get('signature', None): 

43 dims = compiled.meta_['signature'].shape_calculator( 

44 operator.inputs[0].type.shape) 

45 extra_dims = None 

46 else: 

47 N = operator.inputs[0].type.shape[0] 

48 dims = [N] 

49 out = outputs[0] 

50 try: 

51 extra_dims = get_tensor_shape(out.type) 

52 except AttributeError: # pragma: no cover 

53 extra_dims = None 

54 if extra_dims is not None and len(extra_dims) > 0: 

55 operator.outputs[0].shape = list(extra_dims) 

56 else: 

57 operator.outputs[0].type = input_type(dims) 

58 return 

59 

60 if operator.raw_operator.func is not None: 

61 raise TypeError("FunctionTransformer is not supported unless the " 

62 "transform function is of type %r " 

63 "wrapped with onnxnumpy." % type( 

64 operator.raw_operator.func)) 

65 N = operator.inputs[0].type.shape[0] 

66 C = 0 

67 for variable in operator.inputs: 

68 if variable.type.shape[1] is not None: 

69 C += variable.type.shape[1] 

70 else: 

71 C = None 

72 break 

73 

74 operator.outputs[0].type = operator.inputs[0].type.__class__([N, C]) 

75 

76 

77def _copy_attributes(att): 

78 if hasattr(att, 'value'): 

79 return att.value 

80 vt = _var_as_dict(att) 

81 if vt['type']['kind'] == 'tensor': 

82 value = vt['value'] 

83 return make_tensor(att.name, guess_proto_dtype(value.dtype), 

84 value.shape, value.ravel().tolist()) 

85 if vt['type']['kind'] == 'real': 

86 return vt['value'] 

87 raise RuntimeError( # pragma: no cover 

88 f"Unable to copy attribute {att!r}, got {vt!r}.") 

89 

90 

91def new_convert_sklearn_function_transformer(scope, operator, container): 

92 """ 

93 Rewrites the converters implemented in 

94 :epkg:`sklearn-onnx` to support custom functions 

95 implemented with :ref:`l-numpy-onnxpy`. 

96 """ 

97 op = operator.raw_operator 

98 fct = op.func 

99 if hasattr(fct, 'signed_compiled'): 

100 dtype = guess_numpy_type(operator.inputs[0].type) 

101 fct = fct[FctVersion((dtype, ), None)] 

102 if hasattr(fct, 'compiled'): 

103 compiled = fct.compiled 

104 if not hasattr(compiled, 'onnx_'): 

105 raise RuntimeError( # pragma: no cover 

106 "Attribute 'onnx_' is missing, function was not " 

107 "converted to onnx.") 

108 onx = compiled.onnx_ 

109 graph = onx.graph 

110 nodes = graph.node 

111 

112 # renaming all intermediate variables 

113 names = [] 

114 for node in nodes: 

115 for name in node.input: 

116 names.append(name) 

117 for name in node.output: 

118 names.append(name) 

119 names = set(names) 

120 names_mapping = {} 

121 for name in names: 

122 names_mapping[name] = scope.get_unique_variable_name( 

123 f'ft_{name}') 

124 

125 # adding identities 

126 apply_identity(scope, operator.inputs[0].full_name, 

127 names_mapping[graph.input[0].name], container) 

128 apply_identity(scope, names_mapping[graph.output[0].name], 

129 operator.outputs[0].full_name, container) 

130 

131 # adding initializers 

132 for init in graph.initializer: 

133 init = copy.deepcopy(init) 

134 name = names_mapping[init.name] 

135 init.name = name 

136 content = init.SerializeToString() 

137 container.initializers_strings[content] = name 

138 container.initializers.append(init) 

139 

140 # adding nodes 

141 for node in nodes: 

142 atts = {} 

143 for att in node.attribute: 

144 atts[att.name] = _copy_attributes(att) 

145 container.add_node( 

146 node.op_type, 

147 [names_mapping[n] for n in node.input], 

148 [names_mapping[n] for n in node.output], 

149 name=scope.get_unique_operator_name(f'ft_{node.op_type}'), 

150 **atts) 

151 return 

152 

153 if op.func is not None: 

154 raise TypeError( # pragma: no cover 

155 "FunctionTransformer is not supported unless the " 

156 "transform function is of type %r or " 

157 "wrapped with onnxnumpy." % type(op.func)) 

158 if len(operator.inputs) == 1: 

159 apply_identity(scope, operator.inputs[0].full_name, 

160 operator.outputs[0].full_name, container) 

161 else: 

162 apply_concat(scope, [i.full_name for i in operator.inputs], 

163 operator.outputs[0].full_name, container)