1# -*- encoding: utf-8 -*- 

2# pylint: disable=E0203,E1101,C0111 

3""" 

4@file 

5@brief Runtime operator. 

6""" 

7import numpy 

8from ._op import OpRun 

9from ..shape_object import ShapeObject 

10 

11 

12class Scan(OpRun): 

13 

14 atts = { 

15 'body': None, 

16 'num_scan_inputs': None, 

17 'scan_input_axes': [], 

18 'scan_input_directions': [], 

19 'scan_output_axes': [], 

20 'scan_output_directions': [] 

21 } 

22 

23 def __init__(self, onnx_node, desc=None, **options): 

24 OpRun.__init__(self, onnx_node, desc=desc, 

25 expected_attributes=Scan.atts, 

26 **options) 

27 if not hasattr(self.body, 'run'): 

28 raise RuntimeError("Parameter 'body' must have a method 'run', " 

29 "type {}.".format(type(self.body))) 

30 self.input_directions_ = [0 if i >= len(self.scan_input_directions) else self.scan_input_directions[i] 

31 for i in range(self.num_scan_inputs)] 

32 max_dir_in = max(self.input_directions_) 

33 if max_dir_in != 0: 

34 raise RuntimeError( 

35 "Scan is not implemented for other output input_direction than 0.") 

36 self.input_axes_ = [0 if i >= len(self.scan_input_axes) else self.scan_input_axes[i] 

37 for i in range(self.num_scan_inputs)] 

38 max_axe_in = max(self.input_axes_) 

39 if max_axe_in != 0: 

40 raise RuntimeError( 

41 "Scan is not implemented for other input axes than 0.") 

42 self.input_names = self.body.input_names 

43 self.output_names = self.body.output_names 

44 self._run_meth = (self.body.run_in_scan 

45 if hasattr(self.body, 'run_in_scan') 

46 else self.body.run) 

47 

48 def _common_run_shape(self, *args): 

49 num_loop_state_vars = len(args) - self.num_scan_inputs 

50 num_scan_outputs = len(args) - num_loop_state_vars 

51 

52 output_directions = [0 if i >= len(self.scan_output_directions) else self.scan_output_directions[i] 

53 for i in range(num_scan_outputs)] 

54 max_dir_out = max(output_directions) 

55 if max_dir_out != 0: 

56 raise RuntimeError( 

57 "Scan is not implemented for other output output_direction than 0.") 

58 output_axes = [0 if i >= len(self.scan_output_axes) else self.scan_output_axes[i] 

59 for i in range(num_scan_outputs)] 

60 max_axe_out = max(output_axes) 

61 if max_axe_out != 0: 

62 raise RuntimeError( 

63 "Scan is not implemented for other output axes than 0.") 

64 

65 state_names_in = self.input_names[:self.num_scan_inputs] 

66 state_names_out = self.output_names[:len(state_names_in)] 

67 scan_names_in = self.input_names[num_loop_state_vars:] 

68 scan_names_out = self.output_names[num_loop_state_vars:] 

69 scan_values = args[num_loop_state_vars:] 

70 

71 states = args[:num_loop_state_vars] 

72 

73 return (num_loop_state_vars, num_scan_outputs, output_directions, 

74 max_dir_out, output_axes, max_axe_out, state_names_in, 

75 state_names_out, scan_names_in, scan_names_out, 

76 scan_values, states) 

77 

78 def _run(self, *args): # pylint: disable=W0221 

79 (num_loop_state_vars, num_scan_outputs, output_directions, # pylint: disable=W0612 

80 max_dir_out, output_axes, max_axe_out, state_names_in, # pylint: disable=W0612 

81 state_names_out, scan_names_in, scan_names_out, # pylint: disable=W0612 

82 scan_values, states) = self._common_run_shape(*args) # pylint: disable=W0612 

83 

84 max_iter = args[num_loop_state_vars].shape[self.input_axes_[0]] 

85 results = [[] for _ in scan_names_out] 

86 

87 for iter in range(max_iter): 

88 inputs = {} 

89 for name, value in zip(state_names_in, states): 

90 inputs[name] = value 

91 for name, value in zip(scan_names_in, scan_values): 

92 inputs[name] = value[iter] 

93 

94 try: 

95 outputs = self._run_meth(inputs) 

96 except TypeError as e: 

97 raise TypeError( 

98 "Unable to call 'run' for type '{}'.".format( 

99 type(self.body))) from e 

100 

101 states = [outputs[name] for name in state_names_out] 

102 for i, name in enumerate(scan_names_out): 

103 results[i].append(numpy.expand_dims(outputs[name], axis=0)) 

104 

105 for res in results: 

106 conc = numpy.vstack(res) 

107 states.append(conc) 

108 return tuple(states) 

109 

110 def _infer_shapes(self, *args): # pylint: disable=W0221 

111 (num_loop_state_vars, num_scan_outputs, output_directions, # pylint: disable=W0612 

112 max_dir_out, output_axes, max_axe_out, state_names_in, # pylint: disable=W0612 

113 state_names_out, scan_names_in, scan_names_out, # pylint: disable=W0612 

114 scan_values, states) = self._common_run_shape(*args) # pylint: disable=W0612 

115 

116 shapes = list(states) 

117 

118 shape = args[num_loop_state_vars].shape 

119 if shape is None: 

120 for sout in scan_values: 

121 shapes.append(ShapeObject(None, dtype=sout.dtype)) 

122 else: 

123 max_iter = shape[self.input_axes_[0]] 

124 for sout in scan_values: 

125 sc = sout.copy() 

126 sc[0] = max_iter 

127 shapes.append(sc) 

128 

129 return tuple(shapes)