code: def lex_logits(
      lex_h,
      output_projection_w_lex,
      output_projection_b_lex,
      logits_shape,
  ):
      """
      calculate the logits of lexical layer output
      """
      projection_lex_flat = torch.matmul(
          output_projection_w_lex, lex_h.t()
      ).t()
  
      logits = (
          torch.onnx.operators.reshape_from_tensor_shape(
              projection_lex_flat, logits_shape
          )
          + output_projection_b_lex
      )
      return logits
  
label_desc: Use image node_id 2 for calling a global function with example usage: lex_logits(lex_h, output_projection_w_lex, output_projection_b_lex, logits_shape) and returns: logits