code: def _location_from_fx_stack_trace(
      node_stack_trace: str,
  ) -> Optional[diagnostics.infra.Location]:
      """Extract location from FX node stack trace.
  
      TODO(bowbao): Create fx utils module and move this function there.
  
      Args:
          node_stack_trace: The stack trace of the FX node. Example:
  
              File "path/file.py", line 311, in <function>
                  <code>
              |   File "path/file2.py", line 389, in <function>
                  <code>
  
      Returns:
          location: The location of the FX node.
      """
      if "File" not in node_stack_trace:
          return None
  
      lines = node_stack_trace.strip().split("\n")
      idx = 0
      while idx < len(lines) and "File" not in lines[idx]:
          idx += 1
      if idx + 1 >= len(lines):
          return None
  
      pattern = re.compile(r"^File \"(.+)\", line (\d+), in (.+)$")
      matches = pattern.match(lines[idx].strip())
      if matches:
          uri = matches.group(1)
          line_number = int(matches.group(2))
          snippet = lines[idx + 1].strip()
          return diagnostics.infra.Location(
              uri=uri, line=line_number, snippet=snippet
          )
      return None
  
imports: from __future__ import annotations
  import inspect
  import logging
  import operator
  import re
  import types
  from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
  import onnxscript
  from onnxscript.function_libs.torch_lib import graph_building
  import torch
  import torch.fx
  from torch.onnx import _type_utils
  from torch.onnx._internal import _beartype
  from torch.onnx._internal.fx import _pass, diagnostics, onnxfunction_dispatcher, op_validation, type_utils
  from torch.utils import _pytree
label_desc: Use image node_id 3 for calling a global function with example usage: _location_from_fx_stack_trace(node_stack_trace) and returns: None, None, None, diagnostics