code: def jit_model(model, *args) -> Tuple[TinyJit, Dict[int, str]]:
      assert hasattr(model, "forward") or callable(
          model
      ), "model needs a forward function"
  
      @TinyJit
      def run(*x):
          out = (
              model.forward(*x)
              if hasattr(model, "forward")
              else model(*x)
          )
          assert (
              isinstance(out, tuple)
              or isinstance(out, list)
              or isinstance(out, Tensor)
          ), "model output must be a Tensor, tuple, or a list of Tensors for export"
          out = [out] if isinstance(out, Tensor) else out
          return [o.realize() for o in out]
  
      # twice to run the JIT
      for _ in range(2):
          the_output = run(*args)
      special_names = {}
  
      # hack to put the inputs back
      for (j, i), idx in run.input_replace.items():
          realized_input = args[idx].lazydata.base.realized
          run.jit_cache[j].rawbufs[i] = realized_input
          special_names[id(realized_input)] = f"input{idx}"
  
      # TODO: fetch this from the jit in self.input_replace and self.ret (hint: use get_parameters on self.ret)
      for i, output in enumerate(the_output):
          special_names[
              id(output.lazydata.base.realized)
          ] = f"output{i}"
      return run, special_names
  
file_path: tinygrad/extra/export_model.py
imports: from typing import Tuple, Dict, List
  from tinygrad.dtype import DType
  from tinygrad.tensor import Device, Tensor
  from tinygrad.jit import TinyJit
  from tinygrad.nn.state import get_state_dict
  from tinygrad.dtype import dtypes
  import json
label_desc: Use image node_id 2 for calling a global function with example usage: jit_model(model) and returns: run,  special_names, unknown