parameters: (node, **kwargs)
return_statement: return nodes
docstring: Map MXNet's RNN operator attributes to onnx's operators
      and return the created node.
docstring_summary: Map MXNet's RNN operator attributes to onnx's operators
      and return the created node.
docstring_tokens: Map MXNet s RNN operator attributes to onnx s operators and return the created node .
function: def convert_RNN(node, **kwargs):
      """Map MXNet's RNN operator attributes to onnx's operators
      and return the created node.
      """
      from onnx.helper import make_node, make_tensor
      from onnx import TensorProto
  
      name, input_nodes, attrs = get_inputs(node, kwargs)
  
      mode = str(attrs.get('mode'))
      bidirectional = str(attrs.get('bidirectional', 'False'))
      if bidirectional != 'False' and mode not in ['lstm']:
          raise NotImplementedError('Currently RNN onnx export only supports bidirectional is False')
  
      num_layers = int(attrs.get('num_layers', '1'))
  
      use_sequence_length = str(attrs.get('use_sequence_length', 'False'))
      if use_sequence_length != 'False':
          raise NotImplementedError('Currently RNN onnx export only supports use_sequence_length equals to False')
  
      projection_size = str(attrs.get('projection_size', 'None'))
      if projection_size != 'None':
          raise NotImplementedError('Currently RNN onnx export only supports projection_size equals to None')
  
      state_outputs = str(attrs.get('state_outputs', 'False'))
      if state_outputs != 'True':
          raise NotImplementedError('Currently RNN onnx export only supports state_outputs equals to True')
  
      state_size = int(attrs.get('state_size'))
  
      direction = 1
      if bidirectional != 'False':
          direction = 2
  
      data = input_nodes[0]
      param = input_nodes[1]
      dtype = get_input_dtypes(node, kwargs)[2]
  
      create_tensor([0], name+'_0', kwargs['initializer'])
      create_tensor([1], name+'_1', kwargs['initializer'])
      create_tensor([state_size], name+'_state_size', kwargs['initializer'])
      create_tensor([direction], name+'_direction', kwargs['initializer'])
  
      tensor_1 = make_tensor(name+'_1_f', onnx.mapping.NP_TYPE_TO_TENSOR_TYPE[dtype], [1], [1])
  
      nodes = [
          make_node('Shape', [data], [name+'_data_shape']),
          make_node('Split', [name+'_data_shape'], [name+'_seq_length', name+'_batch_size', name+'_input_size']),
          make_node('Concat', [name+'_direction', name+'_batch_size', name+'_state_size'], [name+'_concat'], axis=0),
          make_node('ConstantOfShape', [name+'_concat'], [name+'_COS'], value=tensor_1),
          make_node('Mul', [input_nodes[2], name+'_COS'], [name+'initial_h']),
  
      ]
  
      if mode == 'lstm':
          nodes += [
              make_node('Mul', [input_nodes[3], name+'_COS'], [name+'initial_c']),
          ]
  
          if num_layers == 2:
              if bidirectional != 'False':
                  raise NotImplementedError('Currently lstm onnx export only supports bidirectional when num_layers = 1')
              create_tensor([8*state_size], name+'_8*state_size', kwargs['initializer'])
              create_tensor([4*state_size*state_size], name+'_4*state_size^2', kwargs['initializer'])
              create_tensor([1, 4*state_size, state_size], name+'_WR_shape', kwargs['initializer'])
              create_tensor([1, 8*state_size], name+'_B_shape', kwargs['initializer'])
              create_tensor([4*4*state_size*state_size], name+'_WR_offset', kwargs['initializer'])
  
              nodes += [
                  # Layer 0
                  # get W
                  make_node('Slice', [param, name+'_0', name+'_4*state_size^2'], [name+'_W0_1d']),
                  make_node('Split', [name+'_W0_1d'], [name+'_W00', name+'_W01', name+'_W02', name+'_W03']),
                  make_node('Concat', [name+'_W00', name+'_W03', name+'_W01', name+'_W02'], [name+'_W0_'], axis=0),
                  make_node('Reshape', [name+'_W0_', name+'_WR_shape'], [name+'_W0']),
                  # get R
                  make_node('Add', [name+'_4*state_size^2', name+'_4*state_size^2'], [name+'_R0_offset']),
                  make_node('Slice', [param, name+'_4*state_size^2', name+'_R0_offset'], [name+'_R0_1d']),
                  make_node('Split', [name+'_R0_1d'], [name+'_R00', name+'_R01', name+'_R02', name+'_R03']),
                  make_node('Concat', [name+'_R00', name+'_R03', name+'_R01', name+'_R02'], [name+'_R0_'], axis=0),
                  make_node('Reshape', [name+'_R0_', name+'_WR_shape'], [name+'_R0']),
                  # get B
                  make_node('Add', [name+'_WR_offset', name+'_8*state_size'], [name+'_B0_offset']),
                  make_node('Slice', [param, name+'_WR_offset', name+'_B0_offset'], [name+'_B0_1d']),
                  make_node('Split', [name+'_B0_1d'], [name+'_B00', name+'_B01', name+'_B02', name+'_B03',
                                                       name+'_B04', name+'_B05', name+'_B06', name+'_B07']),
                  make_node('Concat', [name+'_B00', name+'_B03', name+'_B01', name+'_B02',
                                       name+'_B04', name+'_B07', name+'_B05', name+'_B06'], [name+'_B0_'], axis=0),
                  make_node('Reshape', [name+'_B0_', name+'_B_shape'], [name+'_B0']),
                  # get initial states
                  make_node('Split', [name+'initial_h'], [name+'_initial_h0', name+'_initial_h1'], axis=0),
                  make_node('Split', [name+'initial_c'], [name+'_initial_c0', name+'_initial_c1'], axis=0),
                  # get seq_len
                  make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                  make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                  # Layer 0 LSTM
                  make_node('LSTM', [data, name+'_W0', name+'_R0', name+'_B0', name+'_seq_len',
                                     name+'_initial_h0', name+'_initial_c0'],
                            [name+'_lstm0_out_', name+'_lstm0_h', name+'_lstm0_c'], hidden_size=state_size),
                  make_node('Squeeze', [name+'_lstm0_out_'], [name+'_lstm0_out'], axes=[1]),
  
                  # Layer 1
                  # get W
                  make_node('Add', [name+'_R0_offset', name+'_4*state_size^2'], [name+'_W1_offset']),
                  make_node('Slice', [param, name+'_R0_offset', name+'_W1_offset'], [name+'_W1_1d']),
                  make_node('Split', [name+'_W1_1d'], [name+'_W10', name+'_W11', name+'_W12', name+'_W13']),
                  make_node('Concat', [name+'_W10', name+'_W13', name+'_W11', name+'_W12'], [name+'_W1_'], axis=0),
                  make_node('Reshape', [name+'_W1_', name+'_WR_shape'], [name+'_W1']),
                  # get R
                  make_node('Slice', [param, name+'_W1_offset', name+'_WR_offset'], [name+'_R1_1d']),
                  make_node('Split', [name+'_R1_1d'], [name+'_R10', name+'_R11', name+'_R12', name+'_R13']),
                  make_node('Concat', [name+'_R10', name+'_R13', name+'_R11', name+'_R12'], [name+'_R1_'], axis=0),
                  make_node('Reshape', [name+'_R1_', name+'_WR_shape'], [name+'_R1']),
                  # get B
                  make_node('Add', [name+'_B0_offset', name+'_8*state_size'], [name+'_B1_offset']),
                  make_node('Slice', [param, name+'_B0_offset', name+'_B1_offset'], [name+'_B1_1d']),
                  make_node('Split', [name+'_B1_1d'], [name+'_B10', name+'_B11', name+'_B12', name+'_B13',
                                                       name+'_B14', name+'_B15', name+'_B16', name+'_B17']),
                  make_node('Concat', [name+'_B10', name+'_B13', name+'_B11', name+'_B12',
                                       name+'_B14', name+'_B17', name+'_B15', name+'_B16'], [name+'_B1_'], axis=0),
                  make_node('Reshape', [name+'_B1_', name+'_B_shape'], [name+'_B1']),
                  # Layer 1 LSTM
                  make_node('LSTM', [name+'_lstm0_out', name+'_W1', name+'_R1', name+'_B1', name+'_seq_len',
                                     name+'_initial_h1', name+'_initial_c1'],
                            [name+'_lstm1_out_', name+'_lstm1_h', name+'_lstm1_c'], hidden_size=state_size),
                  make_node('Squeeze', [name+'_lstm1_out_'], [name], axes=[1]),
                  make_node('Concat', [name+'_lstm0_h', name+'_lstm1_h'], [name+'1'], axis=0),
                  make_node('Concat', [name+'_lstm0_c', name+'_lstm1_c'], [name+'2'], axis=0),
              ]
          elif num_layers == 1:
              if bidirectional == 'False':
                  create_tensor([4*state_size], name+'_4*state_size', kwargs['initializer'])
                  create_tensor([8*state_size], name+'_8*state_size', kwargs['initializer'])
                  create_tensor([4*state_size*state_size], name+'_4*state_size^2', kwargs['initializer'])
                  create_tensor([1, 4*state_size, state_size], name+'_R_shape', kwargs['initializer'])
                  create_tensor([1, 8*state_size], name+'_B_shape', kwargs['initializer'])
  
                  nodes += [
                      # get W
                      make_node('Mul', [name+'_4*state_size', name+'_input_size'], [name+'_mul0']),
                      make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
                      make_node('Split', [name+'_W_1d'], [name+'_W0', name+'_W1', name+'_W2', name+'_W3']),
                      make_node('Concat', [name+'_W0', name+'_W3', name+'_W1', name+'_W2'], [name+'_W_'], axis=0),
                      make_node('Concat', [name+'_1', name+'_4*state_size', name+'_input_size'],
                                [name+'_W_shape'], axis=0),
                      make_node('Reshape', [name+'_W_', name+'_W_shape'], [name+'_W']),
                      # get R
                      make_node('Add', [name+'_mul0', name+'_4*state_size^2'], [name+'_add0']),
                      make_node('Slice', [param, name+'_mul0', name+'_add0'], [name+'_R_1d']),
                      make_node('Split', [name+'_R_1d'], [name+'_R0', name+'_R1', name+'_R2', name+'_R3']),
                      make_node('Concat', [name+'_R0', name+'_R3', name+'_R1', name+'_R2'], [name+'_R_'], axis=0),
                      make_node('Reshape', [name+'_R_', name+'_R_shape'], [name+'_R']),
                      # get B
                      make_node('Add', [name+'_add0', name+'_8*state_size'], [name+'_add1']),
                      make_node('Slice', [param, name+'_add0', name+'_add1'], [name+'_B_1d']),
                      make_node('Split', [name+'_B_1d'], [name+'_B0', name+'_B1', name+'_B2', name+'_B3',
                                                          name+'_B4', name+'_B5', name+'_B6', name+'_B7']),
                      make_node('Concat', [name+'_B0', name+'_B3', name+'_B1', name+'_B2',
                                           name+'_B4', name+'_B7', name+'_B5', name+'_B6'], [name+'_B_'], axis=0),
                      make_node('Reshape', [name+'_B_', name+'_B_shape'], [name+'_B']),
                      # get seq_len
                      make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                      make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                      # compute LSTM
                      make_node('LSTM', [data, name+'_W', name+'_R', name+'_B',
                                         name+'_seq_len', name+'initial_h', name+'initial_c'],
                                [name+'0_', name+'1', name+'2'], hidden_size=state_size),
                      make_node('Squeeze', [name+'0_'], [name], axes=[1]),
                  ]
              else:
                  create_tensor([-1], name+'_-1', kwargs['initializer'])
                  create_tensor([4*state_size], name+'_4*state_size', kwargs['initializer'])
                  create_tensor([8*state_size], name+'_8*state_size', kwargs['initializer'])
                  create_tensor([4*state_size*state_size], name+'_4*state_size^2', kwargs['initializer'])
                  create_tensor([1, 4*state_size, state_size], name+'_R_shape', kwargs['initializer'])
                  create_tensor([1, 8*state_size], name+'_B_shape', kwargs['initializer'])
  
                  nodes += [
                      # get W_fwd
                      make_node('Mul', [name+'_4*state_size', name+'_input_size'], [name+'_mul0']),
                      make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
                      make_node('Split', [name+'_W_1d'], [name+'_W0', name+'_W1', name+'_W2', name+'_W3']),
                      make_node('Concat', [name+'_W0', name+'_W3', name+'_W1', name+'_W2'],
                                [name+'_W_'], axis=0),
                      make_node('Concat', [name+'_1', name+'_4*state_size', name+'_input_size'],
                                [name+'_W_shape'], axis=0),
                      make_node('Reshape', [name+'_W_', name+'_W_shape'], [name+'_W_fwd']),
                      # get R_fwd
                      make_node('Add', [name+'_mul0', name+'_4*state_size^2'], [name+'_add0']),
                      make_node('Slice', [param, name+'_mul0', name+'_add0'], [name+'_R_1d']),
                      make_node('Split', [name+'_R_1d'], [name+'_R0', name+'_R1', name+'_R2', name+'_R3']),
                      make_node('Concat', [name+'_R0', name+'_R3', name+'_R1', name+'_R2'], [name+'_R_'], axis=0),
                      make_node('Reshape', [name+'_R_', name+'_R_shape'], [name+'_R_fwd']),
                      # get W_bwd
                      make_node('Add', [name+'_add0', name+'_mul0'], [name+'_add1']),
                      make_node('Slice', [param, name+'_add0', name+'_add1'], [name+'_W_1d_bwd']),
                      make_node('Split', [name+'_W_1d_bwd'],
                                [name+'_W0_bwd', name+'_W1_bwd', name+'_W2_bwd', name+'_W3_bwd']),
                      make_node('Concat', [name+'_W0_bwd', name+'_W3_bwd', name+'_W1_bwd', name+'_W2_bwd'],
                                [name+'_W_bwd_'], axis=0),
                      make_node('Reshape', [name+'_W_bwd_', name+'_W_shape'], [name+'_W_bwd']),
                      # get R_bwd
                      make_node('Add', [name+'_add1', name+'_4*state_size^2'], [name+'_add2']),
                      make_node('Slice', [param, name+'_add1', name+'_add2'], [name+'_R_1d_bwd']),
                      make_node('Split', [name+'_R_1d_bwd'],
                                [name+'_R0_bwd', name+'_R1_bwd', name+'_R2_bwd', name+'_R3_bwd']),
                      make_node('Concat', [name+'_R0_bwd', name+'_R3_bwd', name+'_R1_bwd', name+'_R2_bwd'],
                                [name+'_R_bwd_'], axis=0),
                      make_node('Reshape', [name+'_R_bwd_', name+'_R_shape'], [name+'_R_bwd']),
                      # get B_fwd
                      make_node('Add', [name+'_add2', name+'_8*state_size'], [name+'_add3']),
                      make_node('Slice', [param, name+'_add2', name+'_add3'], [name+'_B_1d']),
                      make_node('Split', [name+'_B_1d'], [name+'_B0', name+'_B1', name+'_B2', name+'_B3',
                                                          name+'_B4', name+'_B5', name+'_B6', name+'_B7']),
                      make_node('Concat', [name+'_B0', name+'_B3', name+'_B1', name+'_B2',
                                           name+'_B4', name+'_B7', name+'_B5', name+'_B6'], [name+'_B_'], axis=0),
                      make_node('Reshape', [name+'_B_', name+'_B_shape'], [name+'_B_fwd']),
                      # get B_bwd
                      make_node('Add', [name+'_add3', name+'_8*state_size'], [name+'_add4']),
                      make_node('Slice', [param, name+'_add3', name+'_add4'], [name+'_B_1d_bwd']),
                      make_node('Split', [name+'_B_1d_bwd'],
                                [name+'_B0_bwd', name+'_B1_bwd', name+'_B2_bwd', name+'_B3_bwd',
                                 name+'_B4_bwd', name+'_B5_bwd', name+'_B6_bwd', name+'_B7_bwd']),
                      make_node('Concat', [name+'_B0_bwd', name+'_B3_bwd', name+'_B1_bwd', name+'_B2_bwd',
                                           name+'_B4_bwd', name+'_B7_bwd', name+'_B5_bwd', name+'_B6_bwd'],
                                [name+'_B_bwd_'], axis=0),
                      make_node('Reshape', [name+'_B_bwd_', name+'_B_shape'], [name+'_B_bwd']),
                      # get seq_len
                      make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                      make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                      # compute LSTM
                      make_node('Concat', [name+'_W_fwd', name+'_W_bwd'], [name+'_W'], axis=0),
                      make_node('Concat', [name+'_R_fwd', name+'_R_bwd'], [name+'_R'], axis=0),
                      make_node('Concat', [name+'_B_fwd', name+'_B_bwd'], [name+'_B'], axis=0),
                      make_node('LSTM', [data, name+'_W', name+'_R', name+'_B',
                                         name+'_seq_len', name+'initial_h', name+'initial_c'],
                                [name+'0_', name+'1', name+'2'], hidden_size=state_size, direction='bidirectional'),
                      make_node('Transpose', [name+'0_'], [name+'0_t'], perm=[0, 2, 1, 3]),
                      make_node('Concat', [name+'_seq_length', name+'_batch_size', name+'_-1'],
                                [name+'_shape_out'], axis=0),
                      make_node('Reshape', [name+'0_t', name+'_shape_out'], [name]),
                  ]
          else:
              raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1 or 2')
  
      elif mode == 'gru':
          if num_layers == 2:
              create_tensor([6*state_size], name+'_6*state_size', kwargs['initializer'])
              create_tensor([3*state_size*state_size], name+'_3*state_size^2', kwargs['initializer'])
              create_tensor([1, 3*state_size, state_size], name+'_WR_shape', kwargs['initializer'])
              create_tensor([1, 6*state_size], name+'_B_shape', kwargs['initializer'])
              create_tensor([4*3*state_size*state_size], name+'_WR_offset', kwargs['initializer'])
  
              nodes += [
                  # Layer 0
                  # get W
                  make_node('Slice', [param, name+'_0', name+'_3*state_size^2'], [name+'_W0_1d']),
                  make_node('Split', [name+'_W0_1d'], [name+'_W00', name+'_W01', name+'_W02']),
                  make_node('Concat', [name+'_W01', name+'_W00', name+'_W02'], [name+'_W0_'], axis=0),
                  make_node('Reshape', [name+'_W0_', name+'_WR_shape'], [name+'_W0']),
                  # get R
                  make_node('Add', [name+'_3*state_size^2', name+'_3*state_size^2'], [name+'_R0_offset']),
                  make_node('Slice', [param, name+'_3*state_size^2', name+'_R0_offset'], [name+'_R0_1d']),
                  make_node('Split', [name+'_R0_1d'], [name+'_R00', name+'_R01', name+'_R02']),
                  make_node('Concat', [name+'_R01', name+'_R00', name+'_R02'], [name+'_R0_'], axis=0),
                  make_node('Reshape', [name+'_R0_', name+'_WR_shape'], [name+'_R0']),
                  # get B
                  make_node('Add', [name+'_WR_offset', name+'_6*state_size'], [name+'_B0_offset']),
                  make_node('Slice', [param, name+'_WR_offset', name+'_B0_offset'], [name+'_B0_1d']),
                  make_node('Split', [name+'_B0_1d'], [name+'_B00', name+'_B01', name+'_B02',
                                                       name+'_B03', name+'_B04', name+'_B05']),
                  make_node('Concat', [name+'_B01', name+'_B00', name+'_B02',
                                       name+'_B04', name+'_B03', name+'_B05'], [name+'_B0_'], axis=0),
                  make_node('Reshape', [name+'_B0_', name+'_B_shape'], [name+'_B0']),
                  # get initial states
                  make_node('Split', [name+'initial_h'], [name+'_initial_h0', name+'_initial_h1'], axis=0),
                  # get seq_len
                  make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                  make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                  # Layer 0 GRU
                  make_node('GRU', [data, name+'_W0', name+'_R0', name+'_B0', name+'_seq_len',
                                    name+'_initial_h0'],
                            [name+'_gru0_out_', name+'_gru0_h'], hidden_size=state_size, linear_before_reset=1),
                  make_node('Squeeze', [name+'_gru0_out_'], [name+'_gru0_out'], axes=[1]),
  
                  # Layer 1
                  # get W
                  make_node('Add', [name+'_R0_offset', name+'_3*state_size^2'], [name+'_W1_offset']),
                  make_node('Slice', [param, name+'_R0_offset', name+'_W1_offset'], [name+'_W1_1d']),
                  make_node('Split', [name+'_W1_1d'], [name+'_W10', name+'_W11', name+'_W12']),
                  make_node('Concat', [name+'_W11', name+'_W10', name+'_W12'], [name+'_W1_'], axis=0),
                  make_node('Reshape', [name+'_W1_', name+'_WR_shape'], [name+'_W1']),
                  # get R
                  make_node('Slice', [param, name+'_W1_offset', name+'_WR_offset'], [name+'_R1_1d']),
                  make_node('Split', [name+'_R1_1d'], [name+'_R10', name+'_R11', name+'_R12']),
                  make_node('Concat', [name+'_R11', name+'_R10', name+'_R12'], [name+'_R1_'], axis=0),
                  make_node('Reshape', [name+'_R1_', name+'_WR_shape'], [name+'_R1']),
                  # get B
                  make_node('Add', [name+'_B0_offset', name+'_6*state_size'], [name+'_B1_offset']),
                  make_node('Slice', [param, name+'_B0_offset', name+'_B1_offset'], [name+'_B1_1d']),
                  make_node('Split', [name+'_B1_1d'], [name+'_B10', name+'_B11', name+'_B12',
                                                       name+'_B13', name+'_B14', name+'_B15']),
                  make_node('Concat', [name+'_B11', name+'_B10', name+'_B12',
                                       name+'_B14', name+'_B13', name+'_B15'], [name+'_B1_'], axis=0),
                  make_node('Reshape', [name+'_B1_', name+'_B_shape'], [name+'_B1']),
                  # Layer 1 GRU
                  make_node('GRU', [name+'_gru0_out', name+'_W1', name+'_R1', name+'_B1', name+'_seq_len',
                                    name+'_initial_h1'],
                            [name+'_gru1_out_', name+'_gru1_h'], hidden_size=state_size, linear_before_reset=1),
                  make_node('Squeeze', [name+'_gru1_out_'], [name], axes=[1]),
                  make_node('Concat', [name+'_gru0_h', name+'_gru1_h'], [name+'1'], axis=0)
              ]
  
          elif num_layers == 1:
              create_tensor([3*state_size], name+'_3*state_size', kwargs['initializer'])
              create_tensor([6*state_size], name+'_6*state_size', kwargs['initializer'])
              create_tensor([3*state_size*state_size], name+'_3*state_size^2', kwargs['initializer'])
              create_tensor([1, 3*state_size, state_size], name+'_R_shape', kwargs['initializer'])
              create_tensor([1, 6*state_size], name+'_B_shape', kwargs['initializer'])
  
              nodes += [
                  # get W
                  make_node('Mul', [name+'_3*state_size', name+'_input_size'], [name+'_mul0']),
                  make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
                  make_node('Split', [name+'_W_1d'], [name+'_W0', name+'_W1', name+'_W2']),
                  make_node('Concat', [name+'_W1', name+'_W0', name+'_W2'], [name+'_W_'], axis=0),
                  make_node('Concat', [name+'_1', name+'_3*state_size', name+'_input_size'], [name+'_W_shape'], axis=0),
                  make_node('Reshape', [name+'_W_', name+'_W_shape'], [name+'_W']),
                  # get R
                  make_node('Add', [name+'_mul0', name+'_3*state_size^2'], [name+'_add0']),
                  make_node('Slice', [param, name+'_mul0', name+'_add0'], [name+'_R_1d']),
                  make_node('Split', [name+'_R_1d'], [name+'_R0', name+'_R1', name+'_R2']),
                  make_node('Concat', [name+'_R1', name+'_R0', name+'_R2'], [name+'_R_'], axis=0),
                  make_node('Reshape', [name+'_R_', name+'_R_shape'], [name+'_R']),
                  # get B
                  make_node('Add', [name+'_add0', name+'_6*state_size'], [name+'_add1']),
                  make_node('Slice', [param, name+'_add0', name+'_add1'], [name+'_B_1d']),
                  make_node('Split', [name+'_B_1d'], [name+'_B0', name+'_B1', name+'_B2',
                                                      name+'_B3', name+'_B4', name+'_B5']),
                  make_node('Concat', [name+'_B1', name+'_B0', name+'_B2',
                                       name+'_B4', name+'_B3', name+'_B5'], [name+'_B_'], axis=0),
                  make_node('Reshape', [name+'_B_', name+'_B_shape'], [name+'_B']),
                  # get seq_len
                  make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                  make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                  # compute GRU
                  make_node('GRU', [data, name+'_W', name+'_R', name+'_B', name+'_seq_len', name+'initial_h'],
                            [name+'0_', name+'1'], hidden_size=state_size, linear_before_reset=1),
                  make_node('Squeeze', [name+'0_'], [name], axes=[1]),
              ]
          else:
              raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1 or 2')
  
      elif mode in ['rnn_tanh', 'rnn_relu']:
          activations = ['Tanh']
          if mode == 'rnn_relu':
              activations = ['Relu']
          if num_layers == 2:
              create_tensor([2*state_size], name+'_2*state_size', kwargs['initializer'])
              create_tensor([state_size*state_size], name+'_state_size^2', kwargs['initializer'])
              create_tensor([1, state_size, state_size], name+'_WR_shape', kwargs['initializer'])
              create_tensor([1, 2*state_size], name+'_B_shape', kwargs['initializer'])
              create_tensor([4*state_size*state_size], name+'_WR_offset', kwargs['initializer'])
  
              nodes += [
                  # Layer 0
                  # get W
                  make_node('Slice', [param, name+'_0', name+'_state_size^2'], [name+'_W0_1d']),
                  make_node('Reshape', [name+'_W0_1d', name+'_WR_shape'], [name+'_W0']),
                  # get R
                  make_node('Add', [name+'_state_size^2', name+'_state_size^2'], [name+'_R0_offset']),
                  make_node('Slice', [param, name+'_state_size^2', name+'_R0_offset'], [name+'_R0_1d']),
                  make_node('Reshape', [name+'_R0_1d', name+'_WR_shape'], [name+'_R0']),
                  # get B
                  make_node('Add', [name+'_WR_offset', name+'_2*state_size'], [name+'_B0_offset']),
                  make_node('Slice', [param, name+'_WR_offset', name+'_B0_offset'], [name+'_B0_1d']),
                  make_node('Reshape', [name+'_B0_1d', name+'_B_shape'], [name+'_B0']),
                  # get initial states
                  make_node('Split', [name+'initial_h'], [name+'_initial_h0', name+'_initial_h1'], axis=0),
                  # get seq_len
                  make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                  make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                  # Layer 0 RNN
                  make_node('RNN', [data, name+'_W0', name+'_R0', name+'_B0', name+'_seq_len',
                                    name+'_initial_h0'], [name+'_rnn0_out_', name+'_rnn0_h'],
                            hidden_size=state_size, activations=activations),
                  make_node('Squeeze', [name+'_rnn0_out_'], [name+'_rnn0_out'], axes=[1]),
  
                  # Layer 1
                  # get W
                  make_node('Add', [name+'_R0_offset', name+'_state_size^2'], [name+'_W1_offset']),
                  make_node('Slice', [param, name+'_R0_offset', name+'_W1_offset'], [name+'_W1_1d']),
                  make_node('Reshape', [name+'_W1_1d', name+'_WR_shape'], [name+'_W1']),
                  # get R
                  make_node('Slice', [param, name+'_W1_offset', name+'_WR_offset'], [name+'_R1_1d']),
                  make_node('Reshape', [name+'_R1_1d', name+'_WR_shape'], [name+'_R1']),
                  # get B
                  make_node('Add', [name+'_B0_offset', name+'_2*state_size'], [name+'_B1_offset']),
                  make_node('Slice', [param, name+'_B0_offset', name+'_B1_offset'], [name+'_B1_1d']),
                  make_node('Reshape', [name+'_B1_1d', name+'_B_shape'], [name+'_B1']),
                  # Layer 1 RNN
                  make_node('RNN', [name+'_rnn0_out', name+'_W1', name+'_R1', name+'_B1', name+'_seq_len',
                                    name+'_initial_h1'], [name+'_rnn1_out_', name+'_rnn1_h'],
                            hidden_size=state_size, activations=activations),
                  make_node('Squeeze', [name+'_rnn1_out_'], [name], axes=[1]),
                  make_node('Concat', [name+'_rnn0_h', name+'_rnn1_h'], [name+'1'], axis=0)
              ]
  
          elif num_layers == 1:
              create_tensor([2*state_size], name+'_2*state_size', kwargs['initializer'])
              create_tensor([state_size*state_size], name+'_state_size^2', kwargs['initializer'])
              create_tensor([1, state_size, state_size], name+'_R_shape', kwargs['initializer'])
              create_tensor([1, 2*state_size], name+'_B_shape', kwargs['initializer'])
  
              nodes += [
                  # get W
                  make_node('Mul', [name+'_state_size', name+'_input_size'], [name+'_mul0']),
                  make_node('Slice', [param, name+'_0', name+'_mul0'], [name+'_W_1d']),
                  make_node('Concat', [name+'_1', name+'_state_size', name+'_input_size'], [name+'_W_shape'], axis=0),
                  make_node('Reshape', [name+'_W_1d', name+'_W_shape'], [name+'_W']),
                  # get R
                  make_node('Add', [name+'_mul0', name+'_state_size^2'], [name+'_add0']),
                  make_node('Slice', [param, name+'_mul0', name+'_add0'], [name+'_R_1d']),
                  make_node('Reshape', [name+'_R_1d', name+'_R_shape'], [name+'_R']),
                  # get B
                  make_node('Add', [name+'_add0', name+'_2*state_size'], [name+'_add1']),
                  make_node('Slice', [param, name+'_add0', name+'_add1'], [name+'_B_1d']),
                  make_node('Reshape', [name+'_B_1d', name+'_B_shape'], [name+'_B']),
                  # get seq_len
                  make_node('Tile', [name+'_seq_length', name+'_batch_size'], [name+'_seq_len_']),
                  make_node("Cast", [name+'_seq_len_'], [name+"_seq_len"], to=int(TensorProto.INT32)),
                  # compute RNN
                  make_node('RNN', [data, name+'_W', name+'_R', name+'_B', name+'_seq_len', name+'initial_h'],
                            [name+'0_', name+'1'], hidden_size=state_size, activations=activations),
                  make_node('Squeeze', [name+'0_'], [name], axes=[1]),
              ]
          else:
              raise NotImplementedError('Currently RNN onnx export only supports num_layers equals to 1 or 2')
      else:
          raise NotImplementedError(f"Currently RNN onnx export does not support {mode} mode")
      return nodes
function_tokens: def convert_RNN ( node , * * kwargs ) : from onnx . helper import make_node , make_tensor from onnx import TensorProto name , input_nodes , attrs = get_inputs ( node , kwargs ) mode = str ( attrs . get ( 'mode' ) ) bidirectional = str ( attrs . get ( 'bidirectional' , 'False' ) ) if bidirectional != 'False' and mode not in [ 'lstm' ] : raise NotImplementedError ( 'Currently RNN onnx export only supports bidirectional is False' ) num_layers = int ( attrs . get ( 'num_layers' , '1' ) ) use_sequence_length = str ( attrs . get ( 'use_sequence_length' , 'False' ) ) if use_sequence_length != 'False' : raise NotImplementedError ( 'Currently RNN onnx export only supports use_sequence_length equals to False' ) projection_size = str ( attrs . get ( 'projection_size' , 'None' ) ) if projection_size != 'None' : raise NotImplementedError ( 'Currently RNN onnx export only supports projection_size equals to None' ) state_outputs = str ( attrs . get ( 'state_outputs' , 'False' ) ) if state_outputs != 'True' : raise NotImplementedError ( 'Currently RNN onnx export only supports state_outputs equals to True' ) state_size = int ( attrs . get ( 'state_size' ) ) direction = 1 if bidirectional != 'False' : direction = 2 data = input_nodes [ 0 ] param = input_nodes [ 1 ] dtype = get_input_dtypes ( node , kwargs ) [ 2 ] create_tensor ( [ 0 ] , name + '_0' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 ] , name + '_1' , kwargs [ 'initializer' ] ) create_tensor ( [ state_size ] , name + '_state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ direction ] , name + '_direction' , kwargs [ 'initializer' ] ) tensor_1 = make_tensor ( name + '_1_f' , onnx . mapping . NP_TYPE_TO_TENSOR_TYPE [ dtype ] , [ 1 ] , [ 1 ] ) nodes = [ make_node ( 'Shape' , [ data ] , [ name + '_data_shape' ] ) , make_node ( 'Split' , [ name + '_data_shape' ] , [ name + '_seq_length' , name + '_batch_size' , name + '_input_size' ] ) , make_node ( 'Concat' , [ name + '_direction' , name + '_batch_size' , name + '_state_size' ] , [ name + '_concat' ] , axis = 0 ) , make_node ( 'ConstantOfShape' , [ name + '_concat' ] , [ name + '_COS' ] , value = tensor_1 ) , make_node ( 'Mul' , [ input_nodes [ 2 ] , name + '_COS' ] , [ name + 'initial_h' ] ) , ] if mode == 'lstm' : nodes += [ make_node ( 'Mul' , [ input_nodes [ 3 ] , name + '_COS' ] , [ name + 'initial_c' ] ) , ] if num_layers == 2 : if bidirectional != 'False' : raise NotImplementedError ( 'Currently lstm onnx export only supports bidirectional when num_layers = 1' ) create_tensor ( [ 8 * state_size ] , name + '_8*state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ 4 * state_size * state_size ] , name + '_4*state_size^2' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 4 * state_size , state_size ] , name + '_WR_shape' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 8 * state_size ] , name + '_B_shape' , kwargs [ 'initializer' ] ) create_tensor ( [ 4 * 4 * state_size * state_size ] , name + '_WR_offset' , kwargs [ 'initializer' ] ) nodes += [ # Layer 0 # get W make_node ( 'Slice' , [ param , name + '_0' , name + '_4*state_size^2' ] , [ name + '_W0_1d' ] ) , make_node ( 'Split' , [ name + '_W0_1d' ] , [ name + '_W00' , name + '_W01' , name + '_W02' , name + '_W03' ] ) , make_node ( 'Concat' , [ name + '_W00' , name + '_W03' , name + '_W01' , name + '_W02' ] , [ name + '_W0_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_W0_' , name + '_WR_shape' ] , [ name + '_W0' ] ) , # get R make_node ( 'Add' , [ name + '_4*state_size^2' , name + '_4*state_size^2' ] , [ name + '_R0_offset' ] ) , make_node ( 'Slice' , [ param , name + '_4*state_size^2' , name + '_R0_offset' ] , [ name + '_R0_1d' ] ) , make_node ( 'Split' , [ name + '_R0_1d' ] , [ name + '_R00' , name + '_R01' , name + '_R02' , name + '_R03' ] ) , make_node ( 'Concat' , [ name + '_R00' , name + '_R03' , name + '_R01' , name + '_R02' ] , [ name + '_R0_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_R0_' , name + '_WR_shape' ] , [ name + '_R0' ] ) , # get B make_node ( 'Add' , [ name + '_WR_offset' , name + '_8*state_size' ] , [ name + '_B0_offset' ] ) , make_node ( 'Slice' , [ param , name + '_WR_offset' , name + '_B0_offset' ] , [ name + '_B0_1d' ] ) , make_node ( 'Split' , [ name + '_B0_1d' ] , [ name + '_B00' , name + '_B01' , name + '_B02' , name + '_B03' , name + '_B04' , name + '_B05' , name + '_B06' , name + '_B07' ] ) , make_node ( 'Concat' , [ name + '_B00' , name + '_B03' , name + '_B01' , name + '_B02' , name + '_B04' , name + '_B07' , name + '_B05' , name + '_B06' ] , [ name + '_B0_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_B0_' , name + '_B_shape' ] , [ name + '_B0' ] ) , # get initial states make_node ( 'Split' , [ name + 'initial_h' ] , [ name + '_initial_h0' , name + '_initial_h1' ] , axis = 0 ) , make_node ( 'Split' , [ name + 'initial_c' ] , [ name + '_initial_c0' , name + '_initial_c1' ] , axis = 0 ) , # get seq_len make_node ( 'Tile' , [ name + '_seq_length' , name + '_batch_size' ] , [ name + '_seq_len_' ] ) , make_node ( "Cast" , [ name + '_seq_len_' ] , [ name + "_seq_len" ] , to = int ( TensorProto . INT32 ) ) , # Layer 0 LSTM make_node ( 'LSTM' , [ data , name + '_W0' , name + '_R0' , name + '_B0' , name + '_seq_len' , name + '_initial_h0' , name + '_initial_c0' ] , [ name + '_lstm0_out_' , name + '_lstm0_h' , name + '_lstm0_c' ] , hidden_size = state_size ) , make_node ( 'Squeeze' , [ name + '_lstm0_out_' ] , [ name + '_lstm0_out' ] , axes = [ 1 ] ) , # Layer 1 # get W make_node ( 'Add' , [ name + '_R0_offset' , name + '_4*state_size^2' ] , [ name + '_W1_offset' ] ) , make_node ( 'Slice' , [ param , name + '_R0_offset' , name + '_W1_offset' ] , [ name + '_W1_1d' ] ) , make_node ( 'Split' , [ name + '_W1_1d' ] , [ name + '_W10' , name + '_W11' , name + '_W12' , name + '_W13' ] ) , make_node ( 'Concat' , [ name + '_W10' , name + '_W13' , name + '_W11' , name + '_W12' ] , [ name + '_W1_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_W1_' , name + '_WR_shape' ] , [ name + '_W1' ] ) , # get R make_node ( 'Slice' , [ param , name + '_W1_offset' , name + '_WR_offset' ] , [ name + '_R1_1d' ] ) , make_node ( 'Split' , [ name + '_R1_1d' ] , [ name + '_R10' , name + '_R11' , name + '_R12' , name + '_R13' ] ) , make_node ( 'Concat' , [ name + '_R10' , name + '_R13' , name + '_R11' , name + '_R12' ] , [ name + '_R1_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_R1_' , name + '_WR_shape' ] , [ name + '_R1' ] ) , # get B make_node ( 'Add' , [ name + '_B0_offset' , name + '_8*state_size' ] , [ name + '_B1_offset' ] ) , make_node ( 'Slice' , [ param , name + '_B0_offset' , name + '_B1_offset' ] , [ name + '_B1_1d' ] ) , make_node ( 'Split' , [ name + '_B1_1d' ] , [ name + '_B10' , name + '_B11' , name + '_B12' , name + '_B13' , name + '_B14' , name + '_B15' , name + '_B16' , name + '_B17' ] ) , make_node ( 'Concat' , [ name + '_B10' , name + '_B13' , name + '_B11' , name + '_B12' , name + '_B14' , name + '_B17' , name + '_B15' , name + '_B16' ] , [ name + '_B1_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_B1_' , name + '_B_shape' ] , [ name + '_B1' ] ) , # Layer 1 LSTM make_node ( 'LSTM' , [ name + '_lstm0_out' , name + '_W1' , name + '_R1' , name + '_B1' , name + '_seq_len' , name + '_initial_h1' , name + '_initial_c1' ] , [ name + '_lstm1_out_' , name + '_lstm1_h' , name + '_lstm1_c' ] , hidden_size = state_size ) , make_node ( 'Squeeze' , [ name + '_lstm1_out_' ] , [ name ] , axes = [ 1 ] ) , make_node ( 'Concat' , [ name + '_lstm0_h' , name + '_lstm1_h' ] , [ name + '1' ] , axis = 0 ) , make_node ( 'Concat' , [ name + '_lstm0_c' , name + '_lstm1_c' ] , [ name + '2' ] , axis = 0 ) , ] elif num_layers == 1 : if bidirectional == 'False' : create_tensor ( [ 4 * state_size ] , name + '_4*state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ 8 * state_size ] , name + '_8*state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ 4 * state_size * state_size ] , name + '_4*state_size^2' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 4 * state_size , state_size ] , name + '_R_shape' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 8 * state_size ] , name + '_B_shape' , kwargs [ 'initializer' ] ) nodes += [ # get W make_node ( 'Mul' , [ name + '_4*state_size' , name + '_input_size' ] , [ name + '_mul0' ] ) , make_node ( 'Slice' , [ param , name + '_0' , name + '_mul0' ] , [ name + '_W_1d' ] ) , make_node ( 'Split' , [ name + '_W_1d' ] , [ name + '_W0' , name + '_W1' , name + '_W2' , name + '_W3' ] ) , make_node ( 'Concat' , [ name + '_W0' , name + '_W3' , name + '_W1' , name + '_W2' ] , [ name + '_W_' ] , axis = 0 ) , make_node ( 'Concat' , [ name + '_1' , name + '_4*state_size' , name + '_input_size' ] , [ name + '_W_shape' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_W_' , name + '_W_shape' ] , [ name + '_W' ] ) , # get R make_node ( 'Add' , [ name + '_mul0' , name + '_4*state_size^2' ] , [ name + '_add0' ] ) , make_node ( 'Slice' , [ param , name + '_mul0' , name + '_add0' ] , [ name + '_R_1d' ] ) , make_node ( 'Split' , [ name + '_R_1d' ] , [ name + '_R0' , name + '_R1' , name + '_R2' , name + '_R3' ] ) , make_node ( 'Concat' , [ name + '_R0' , name + '_R3' , name + '_R1' , name + '_R2' ] , [ name + '_R_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_R_' , name + '_R_shape' ] , [ name + '_R' ] ) , # get B make_node ( 'Add' , [ name + '_add0' , name + '_8*state_size' ] , [ name + '_add1' ] ) , make_node ( 'Slice' , [ param , name + '_add0' , name + '_add1' ] , [ name + '_B_1d' ] ) , make_node ( 'Split' , [ name + '_B_1d' ] , [ name + '_B0' , name + '_B1' , name + '_B2' , name + '_B3' , name + '_B4' , name + '_B5' , name + '_B6' , name + '_B7' ] ) , make_node ( 'Concat' , [ name + '_B0' , name + '_B3' , name + '_B1' , name + '_B2' , name + '_B4' , name + '_B7' , name + '_B5' , name + '_B6' ] , [ name + '_B_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_B_' , name + '_B_shape' ] , [ name + '_B' ] ) , # get seq_len make_node ( 'Tile' , [ name + '_seq_length' , name + '_batch_size' ] , [ name + '_seq_len_' ] ) , make_node ( "Cast" , [ name + '_seq_len_' ] , [ name + "_seq_len" ] , to = int ( TensorProto . INT32 ) ) , # compute LSTM make_node ( 'LSTM' , [ data , name + '_W' , name + '_R' , name + '_B' , name + '_seq_len' , name + 'initial_h' , name + 'initial_c' ] , [ name + '0_' , name + '1' , name + '2' ] , hidden_size = state_size ) , make_node ( 'Squeeze' , [ name + '0_' ] , [ name ] , axes = [ 1 ] ) , ] else : create_tensor ( [ - 1 ] , name + '_-1' , kwargs [ 'initializer' ] ) create_tensor ( [ 4 * state_size ] , name + '_4*state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ 8 * state_size ] , name + '_8*state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ 4 * state_size * state_size ] , name + '_4*state_size^2' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 4 * state_size , state_size ] , name + '_R_shape' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 8 * state_size ] , name + '_B_shape' , kwargs [ 'initializer' ] ) nodes += [ # get W_fwd make_node ( 'Mul' , [ name + '_4*state_size' , name + '_input_size' ] , [ name + '_mul0' ] ) , make_node ( 'Slice' , [ param , name + '_0' , name + '_mul0' ] , [ name + '_W_1d' ] ) , make_node ( 'Split' , [ name + '_W_1d' ] , [ name + '_W0' , name + '_W1' , name + '_W2' , name + '_W3' ] ) , make_node ( 'Concat' , [ name + '_W0' , name + '_W3' , name + '_W1' , name + '_W2' ] , [ name + '_W_' ] , axis = 0 ) , make_node ( 'Concat' , [ name + '_1' , name + '_4*state_size' , name + '_input_size' ] , [ name + '_W_shape' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_W_' , name + '_W_shape' ] , [ name + '_W_fwd' ] ) , # get R_fwd make_node ( 'Add' , [ name + '_mul0' , name + '_4*state_size^2' ] , [ name + '_add0' ] ) , make_node ( 'Slice' , [ param , name + '_mul0' , name + '_add0' ] , [ name + '_R_1d' ] ) , make_node ( 'Split' , [ name + '_R_1d' ] , [ name + '_R0' , name + '_R1' , name + '_R2' , name + '_R3' ] ) , make_node ( 'Concat' , [ name + '_R0' , name + '_R3' , name + '_R1' , name + '_R2' ] , [ name + '_R_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_R_' , name + '_R_shape' ] , [ name + '_R_fwd' ] ) , # get W_bwd make_node ( 'Add' , [ name + '_add0' , name + '_mul0' ] , [ name + '_add1' ] ) , make_node ( 'Slice' , [ param , name + '_add0' , name + '_add1' ] , [ name + '_W_1d_bwd' ] ) , make_node ( 'Split' , [ name + '_W_1d_bwd' ] , [ name + '_W0_bwd' , name + '_W1_bwd' , name + '_W2_bwd' , name + '_W3_bwd' ] ) , make_node ( 'Concat' , [ name + '_W0_bwd' , name + '_W3_bwd' , name + '_W1_bwd' , name + '_W2_bwd' ] , [ name + '_W_bwd_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_W_bwd_' , name + '_W_shape' ] , [ name + '_W_bwd' ] ) , # get R_bwd make_node ( 'Add' , [ name + '_add1' , name + '_4*state_size^2' ] , [ name + '_add2' ] ) , make_node ( 'Slice' , [ param , name + '_add1' , name + '_add2' ] , [ name + '_R_1d_bwd' ] ) , make_node ( 'Split' , [ name + '_R_1d_bwd' ] , [ name + '_R0_bwd' , name + '_R1_bwd' , name + '_R2_bwd' , name + '_R3_bwd' ] ) , make_node ( 'Concat' , [ name + '_R0_bwd' , name + '_R3_bwd' , name + '_R1_bwd' , name + '_R2_bwd' ] , [ name + '_R_bwd_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_R_bwd_' , name + '_R_shape' ] , [ name + '_R_bwd' ] ) , # get B_fwd make_node ( 'Add' , [ name + '_add2' , name + '_8*state_size' ] , [ name + '_add3' ] ) , make_node ( 'Slice' , [ param , name + '_add2' , name + '_add3' ] , [ name + '_B_1d' ] ) , make_node ( 'Split' , [ name + '_B_1d' ] , [ name + '_B0' , name + '_B1' , name + '_B2' , name + '_B3' , name + '_B4' , name + '_B5' , name + '_B6' , name + '_B7' ] ) , make_node ( 'Concat' , [ name + '_B0' , name + '_B3' , name + '_B1' , name + '_B2' , name + '_B4' , name + '_B7' , name + '_B5' , name + '_B6' ] , [ name + '_B_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_B_' , name + '_B_shape' ] , [ name + '_B_fwd' ] ) , # get B_bwd make_node ( 'Add' , [ name + '_add3' , name + '_8*state_size' ] , [ name + '_add4' ] ) , make_node ( 'Slice' , [ param , name + '_add3' , name + '_add4' ] , [ name + '_B_1d_bwd' ] ) , make_node ( 'Split' , [ name + '_B_1d_bwd' ] , [ name + '_B0_bwd' , name + '_B1_bwd' , name + '_B2_bwd' , name + '_B3_bwd' , name + '_B4_bwd' , name + '_B5_bwd' , name + '_B6_bwd' , name + '_B7_bwd' ] ) , make_node ( 'Concat' , [ name + '_B0_bwd' , name + '_B3_bwd' , name + '_B1_bwd' , name + '_B2_bwd' , name + '_B4_bwd' , name + '_B7_bwd' , name + '_B5_bwd' , name + '_B6_bwd' ] , [ name + '_B_bwd_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_B_bwd_' , name + '_B_shape' ] , [ name + '_B_bwd' ] ) , # get seq_len make_node ( 'Tile' , [ name + '_seq_length' , name + '_batch_size' ] , [ name + '_seq_len_' ] ) , make_node ( "Cast" , [ name + '_seq_len_' ] , [ name + "_seq_len" ] , to = int ( TensorProto . INT32 ) ) , # compute LSTM make_node ( 'Concat' , [ name + '_W_fwd' , name + '_W_bwd' ] , [ name + '_W' ] , axis = 0 ) , make_node ( 'Concat' , [ name + '_R_fwd' , name + '_R_bwd' ] , [ name + '_R' ] , axis = 0 ) , make_node ( 'Concat' , [ name + '_B_fwd' , name + '_B_bwd' ] , [ name + '_B' ] , axis = 0 ) , make_node ( 'LSTM' , [ data , name + '_W' , name + '_R' , name + '_B' , name + '_seq_len' , name + 'initial_h' , name + 'initial_c' ] , [ name + '0_' , name + '1' , name + '2' ] , hidden_size = state_size , direction = 'bidirectional' ) , make_node ( 'Transpose' , [ name + '0_' ] , [ name + '0_t' ] , perm = [ 0 , 2 , 1 , 3 ] ) , make_node ( 'Concat' , [ name + '_seq_length' , name + '_batch_size' , name + '_-1' ] , [ name + '_shape_out' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '0_t' , name + '_shape_out' ] , [ name ] ) , ] else : raise NotImplementedError ( 'Currently RNN onnx export only supports num_layers equals to 1 or 2' ) elif mode == 'gru' : if num_layers == 2 : create_tensor ( [ 6 * state_size ] , name + '_6*state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ 3 * state_size * state_size ] , name + '_3*state_size^2' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 3 * state_size , state_size ] , name + '_WR_shape' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 6 * state_size ] , name + '_B_shape' , kwargs [ 'initializer' ] ) create_tensor ( [ 4 * 3 * state_size * state_size ] , name + '_WR_offset' , kwargs [ 'initializer' ] ) nodes += [ # Layer 0 # get W make_node ( 'Slice' , [ param , name + '_0' , name + '_3*state_size^2' ] , [ name + '_W0_1d' ] ) , make_node ( 'Split' , [ name + '_W0_1d' ] , [ name + '_W00' , name + '_W01' , name + '_W02' ] ) , make_node ( 'Concat' , [ name + '_W01' , name + '_W00' , name + '_W02' ] , [ name + '_W0_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_W0_' , name + '_WR_shape' ] , [ name + '_W0' ] ) , # get R make_node ( 'Add' , [ name + '_3*state_size^2' , name + '_3*state_size^2' ] , [ name + '_R0_offset' ] ) , make_node ( 'Slice' , [ param , name + '_3*state_size^2' , name + '_R0_offset' ] , [ name + '_R0_1d' ] ) , make_node ( 'Split' , [ name + '_R0_1d' ] , [ name + '_R00' , name + '_R01' , name + '_R02' ] ) , make_node ( 'Concat' , [ name + '_R01' , name + '_R00' , name + '_R02' ] , [ name + '_R0_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_R0_' , name + '_WR_shape' ] , [ name + '_R0' ] ) , # get B make_node ( 'Add' , [ name + '_WR_offset' , name + '_6*state_size' ] , [ name + '_B0_offset' ] ) , make_node ( 'Slice' , [ param , name + '_WR_offset' , name + '_B0_offset' ] , [ name + '_B0_1d' ] ) , make_node ( 'Split' , [ name + '_B0_1d' ] , [ name + '_B00' , name + '_B01' , name + '_B02' , name + '_B03' , name + '_B04' , name + '_B05' ] ) , make_node ( 'Concat' , [ name + '_B01' , name + '_B00' , name + '_B02' , name + '_B04' , name + '_B03' , name + '_B05' ] , [ name + '_B0_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_B0_' , name + '_B_shape' ] , [ name + '_B0' ] ) , # get initial states make_node ( 'Split' , [ name + 'initial_h' ] , [ name + '_initial_h0' , name + '_initial_h1' ] , axis = 0 ) , # get seq_len make_node ( 'Tile' , [ name + '_seq_length' , name + '_batch_size' ] , [ name + '_seq_len_' ] ) , make_node ( "Cast" , [ name + '_seq_len_' ] , [ name + "_seq_len" ] , to = int ( TensorProto . INT32 ) ) , # Layer 0 GRU make_node ( 'GRU' , [ data , name + '_W0' , name + '_R0' , name + '_B0' , name + '_seq_len' , name + '_initial_h0' ] , [ name + '_gru0_out_' , name + '_gru0_h' ] , hidden_size = state_size , linear_before_reset = 1 ) , make_node ( 'Squeeze' , [ name + '_gru0_out_' ] , [ name + '_gru0_out' ] , axes = [ 1 ] ) , # Layer 1 # get W make_node ( 'Add' , [ name + '_R0_offset' , name + '_3*state_size^2' ] , [ name + '_W1_offset' ] ) , make_node ( 'Slice' , [ param , name + '_R0_offset' , name + '_W1_offset' ] , [ name + '_W1_1d' ] ) , make_node ( 'Split' , [ name + '_W1_1d' ] , [ name + '_W10' , name + '_W11' , name + '_W12' ] ) , make_node ( 'Concat' , [ name + '_W11' , name + '_W10' , name + '_W12' ] , [ name + '_W1_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_W1_' , name + '_WR_shape' ] , [ name + '_W1' ] ) , # get R make_node ( 'Slice' , [ param , name + '_W1_offset' , name + '_WR_offset' ] , [ name + '_R1_1d' ] ) , make_node ( 'Split' , [ name + '_R1_1d' ] , [ name + '_R10' , name + '_R11' , name + '_R12' ] ) , make_node ( 'Concat' , [ name + '_R11' , name + '_R10' , name + '_R12' ] , [ name + '_R1_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_R1_' , name + '_WR_shape' ] , [ name + '_R1' ] ) , # get B make_node ( 'Add' , [ name + '_B0_offset' , name + '_6*state_size' ] , [ name + '_B1_offset' ] ) , make_node ( 'Slice' , [ param , name + '_B0_offset' , name + '_B1_offset' ] , [ name + '_B1_1d' ] ) , make_node ( 'Split' , [ name + '_B1_1d' ] , [ name + '_B10' , name + '_B11' , name + '_B12' , name + '_B13' , name + '_B14' , name + '_B15' ] ) , make_node ( 'Concat' , [ name + '_B11' , name + '_B10' , name + '_B12' , name + '_B14' , name + '_B13' , name + '_B15' ] , [ name + '_B1_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_B1_' , name + '_B_shape' ] , [ name + '_B1' ] ) , # Layer 1 GRU make_node ( 'GRU' , [ name + '_gru0_out' , name + '_W1' , name + '_R1' , name + '_B1' , name + '_seq_len' , name + '_initial_h1' ] , [ name + '_gru1_out_' , name + '_gru1_h' ] , hidden_size = state_size , linear_before_reset = 1 ) , make_node ( 'Squeeze' , [ name + '_gru1_out_' ] , [ name ] , axes = [ 1 ] ) , make_node ( 'Concat' , [ name + '_gru0_h' , name + '_gru1_h' ] , [ name + '1' ] , axis = 0 ) ] elif num_layers == 1 : create_tensor ( [ 3 * state_size ] , name + '_3*state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ 6 * state_size ] , name + '_6*state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ 3 * state_size * state_size ] , name + '_3*state_size^2' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 3 * state_size , state_size ] , name + '_R_shape' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 6 * state_size ] , name + '_B_shape' , kwargs [ 'initializer' ] ) nodes += [ # get W make_node ( 'Mul' , [ name + '_3*state_size' , name + '_input_size' ] , [ name + '_mul0' ] ) , make_node ( 'Slice' , [ param , name + '_0' , name + '_mul0' ] , [ name + '_W_1d' ] ) , make_node ( 'Split' , [ name + '_W_1d' ] , [ name + '_W0' , name + '_W1' , name + '_W2' ] ) , make_node ( 'Concat' , [ name + '_W1' , name + '_W0' , name + '_W2' ] , [ name + '_W_' ] , axis = 0 ) , make_node ( 'Concat' , [ name + '_1' , name + '_3*state_size' , name + '_input_size' ] , [ name + '_W_shape' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_W_' , name + '_W_shape' ] , [ name + '_W' ] ) , # get R make_node ( 'Add' , [ name + '_mul0' , name + '_3*state_size^2' ] , [ name + '_add0' ] ) , make_node ( 'Slice' , [ param , name + '_mul0' , name + '_add0' ] , [ name + '_R_1d' ] ) , make_node ( 'Split' , [ name + '_R_1d' ] , [ name + '_R0' , name + '_R1' , name + '_R2' ] ) , make_node ( 'Concat' , [ name + '_R1' , name + '_R0' , name + '_R2' ] , [ name + '_R_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_R_' , name + '_R_shape' ] , [ name + '_R' ] ) , # get B make_node ( 'Add' , [ name + '_add0' , name + '_6*state_size' ] , [ name + '_add1' ] ) , make_node ( 'Slice' , [ param , name + '_add0' , name + '_add1' ] , [ name + '_B_1d' ] ) , make_node ( 'Split' , [ name + '_B_1d' ] , [ name + '_B0' , name + '_B1' , name + '_B2' , name + '_B3' , name + '_B4' , name + '_B5' ] ) , make_node ( 'Concat' , [ name + '_B1' , name + '_B0' , name + '_B2' , name + '_B4' , name + '_B3' , name + '_B5' ] , [ name + '_B_' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_B_' , name + '_B_shape' ] , [ name + '_B' ] ) , # get seq_len make_node ( 'Tile' , [ name + '_seq_length' , name + '_batch_size' ] , [ name + '_seq_len_' ] ) , make_node ( "Cast" , [ name + '_seq_len_' ] , [ name + "_seq_len" ] , to = int ( TensorProto . INT32 ) ) , # compute GRU make_node ( 'GRU' , [ data , name + '_W' , name + '_R' , name + '_B' , name + '_seq_len' , name + 'initial_h' ] , [ name + '0_' , name + '1' ] , hidden_size = state_size , linear_before_reset = 1 ) , make_node ( 'Squeeze' , [ name + '0_' ] , [ name ] , axes = [ 1 ] ) , ] else : raise NotImplementedError ( 'Currently RNN onnx export only supports num_layers equals to 1 or 2' ) elif mode in [ 'rnn_tanh' , 'rnn_relu' ] : activations = [ 'Tanh' ] if mode == 'rnn_relu' : activations = [ 'Relu' ] if num_layers == 2 : create_tensor ( [ 2 * state_size ] , name + '_2*state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ state_size * state_size ] , name + '_state_size^2' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , state_size , state_size ] , name + '_WR_shape' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 2 * state_size ] , name + '_B_shape' , kwargs [ 'initializer' ] ) create_tensor ( [ 4 * state_size * state_size ] , name + '_WR_offset' , kwargs [ 'initializer' ] ) nodes += [ # Layer 0 # get W make_node ( 'Slice' , [ param , name + '_0' , name + '_state_size^2' ] , [ name + '_W0_1d' ] ) , make_node ( 'Reshape' , [ name + '_W0_1d' , name + '_WR_shape' ] , [ name + '_W0' ] ) , # get R make_node ( 'Add' , [ name + '_state_size^2' , name + '_state_size^2' ] , [ name + '_R0_offset' ] ) , make_node ( 'Slice' , [ param , name + '_state_size^2' , name + '_R0_offset' ] , [ name + '_R0_1d' ] ) , make_node ( 'Reshape' , [ name + '_R0_1d' , name + '_WR_shape' ] , [ name + '_R0' ] ) , # get B make_node ( 'Add' , [ name + '_WR_offset' , name + '_2*state_size' ] , [ name + '_B0_offset' ] ) , make_node ( 'Slice' , [ param , name + '_WR_offset' , name + '_B0_offset' ] , [ name + '_B0_1d' ] ) , make_node ( 'Reshape' , [ name + '_B0_1d' , name + '_B_shape' ] , [ name + '_B0' ] ) , # get initial states make_node ( 'Split' , [ name + 'initial_h' ] , [ name + '_initial_h0' , name + '_initial_h1' ] , axis = 0 ) , # get seq_len make_node ( 'Tile' , [ name + '_seq_length' , name + '_batch_size' ] , [ name + '_seq_len_' ] ) , make_node ( "Cast" , [ name + '_seq_len_' ] , [ name + "_seq_len" ] , to = int ( TensorProto . INT32 ) ) , # Layer 0 RNN make_node ( 'RNN' , [ data , name + '_W0' , name + '_R0' , name + '_B0' , name + '_seq_len' , name + '_initial_h0' ] , [ name + '_rnn0_out_' , name + '_rnn0_h' ] , hidden_size = state_size , activations = activations ) , make_node ( 'Squeeze' , [ name + '_rnn0_out_' ] , [ name + '_rnn0_out' ] , axes = [ 1 ] ) , # Layer 1 # get W make_node ( 'Add' , [ name + '_R0_offset' , name + '_state_size^2' ] , [ name + '_W1_offset' ] ) , make_node ( 'Slice' , [ param , name + '_R0_offset' , name + '_W1_offset' ] , [ name + '_W1_1d' ] ) , make_node ( 'Reshape' , [ name + '_W1_1d' , name + '_WR_shape' ] , [ name + '_W1' ] ) , # get R make_node ( 'Slice' , [ param , name + '_W1_offset' , name + '_WR_offset' ] , [ name + '_R1_1d' ] ) , make_node ( 'Reshape' , [ name + '_R1_1d' , name + '_WR_shape' ] , [ name + '_R1' ] ) , # get B make_node ( 'Add' , [ name + '_B0_offset' , name + '_2*state_size' ] , [ name + '_B1_offset' ] ) , make_node ( 'Slice' , [ param , name + '_B0_offset' , name + '_B1_offset' ] , [ name + '_B1_1d' ] ) , make_node ( 'Reshape' , [ name + '_B1_1d' , name + '_B_shape' ] , [ name + '_B1' ] ) , # Layer 1 RNN make_node ( 'RNN' , [ name + '_rnn0_out' , name + '_W1' , name + '_R1' , name + '_B1' , name + '_seq_len' , name + '_initial_h1' ] , [ name + '_rnn1_out_' , name + '_rnn1_h' ] , hidden_size = state_size , activations = activations ) , make_node ( 'Squeeze' , [ name + '_rnn1_out_' ] , [ name ] , axes = [ 1 ] ) , make_node ( 'Concat' , [ name + '_rnn0_h' , name + '_rnn1_h' ] , [ name + '1' ] , axis = 0 ) ] elif num_layers == 1 : create_tensor ( [ 2 * state_size ] , name + '_2*state_size' , kwargs [ 'initializer' ] ) create_tensor ( [ state_size * state_size ] , name + '_state_size^2' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , state_size , state_size ] , name + '_R_shape' , kwargs [ 'initializer' ] ) create_tensor ( [ 1 , 2 * state_size ] , name + '_B_shape' , kwargs [ 'initializer' ] ) nodes += [ # get W make_node ( 'Mul' , [ name + '_state_size' , name + '_input_size' ] , [ name + '_mul0' ] ) , make_node ( 'Slice' , [ param , name + '_0' , name + '_mul0' ] , [ name + '_W_1d' ] ) , make_node ( 'Concat' , [ name + '_1' , name + '_state_size' , name + '_input_size' ] , [ name + '_W_shape' ] , axis = 0 ) , make_node ( 'Reshape' , [ name + '_W_1d' , name + '_W_shape' ] , [ name + '_W' ] ) , # get R make_node ( 'Add' , [ name + '_mul0' , name + '_state_size^2' ] , [ name + '_add0' ] ) , make_node ( 'Slice' , [ param , name + '_mul0' , name + '_add0' ] , [ name + '_R_1d' ] ) , make_node ( 'Reshape' , [ name + '_R_1d' , name + '_R_shape' ] , [ name + '_R' ] ) , # get B make_node ( 'Add' , [ name + '_add0' , name + '_2*state_size' ] , [ name + '_add1' ] ) , make_node ( 'Slice' , [ param , name + '_add0' , name + '_add1' ] , [ name + '_B_1d' ] ) , make_node ( 'Reshape' , [ name + '_B_1d' , name + '_B_shape' ] , [ name + '_B' ] ) , # get seq_len make_node ( 'Tile' , [ name + '_seq_length' , name + '_batch_size' ] , [ name + '_seq_len_' ] ) , make_node ( "Cast" , [ name + '_seq_len_' ] , [ name + "_seq_len" ] , to = int ( TensorProto . INT32 ) ) , # compute RNN make_node ( 'RNN' , [ data , name + '_W' , name + '_R' , name + '_B' , name + '_seq_len' , name + 'initial_h' ] , [ name + '0_' , name + '1' ] , hidden_size = state_size , activations = activations ) , make_node ( 'Squeeze' , [ name + '0_' ] , [ name ] , axes = [ 1 ] ) , ] else : raise NotImplementedError ( 'Currently RNN onnx export only supports num_layers equals to 1 or 2' ) else : raise NotImplementedError ( f"Currently RNN onnx export does not support {mode} mode" ) return nodes