code: def export_model_webgpu(
      functions,
      statements,
      bufs,
      bufs_to_save,
      weight_names,
      input_names,
      output_names,
  ) -> Tuple[str, int, int]:
      kernel_code = "\n\n".join(
          [
              f"const {key} = `{code.replace(key, 'main')}`;"
              for key, code in functions.items()
          ]
      )
      kernel_names = ", ".join(
          [
              name
              for (name, _args, _global_size, _local_size) in statements
          ]
      )
      kernel_calls = "\n        ".join(
          [
              f"addComputePass(device, commandEncoder, piplines[{i}], [{', '.join(args)}], {global_size});"
              for i, (
                  _name,
                  args,
                  global_size,
                  _local_size,
              ) in enumerate(statements)
          ]
      )
      _bufs = "\n    ".join(
          [
              f"const {name} = "
              + (
                  f"createEmptyBuf(device, {size});"
                  if _key not in weight_names
                  else f"createWeightBuf(device, {size}, getTensorBuffer(safetensor, metadata['{weight_names[_key]}']))"
              )
              + ";"
              for name, (size, dtype, _key) in bufs.items()
          ]
      )
      gpu_write_bufs = "\n    ".join(
          [
              f"const gpuWriteBuffer{i} = device.createBuffer({{size:{input_name}.size, usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.MAP_WRITE }});"
              for i, input_name in enumerate(input_names)
          ]
      )
      input_writers = "\n    ".join(
          [
              f"await gpuWriteBuffer{i}.mapAsync(GPUMapMode.WRITE);\n        new Float32Array(gpuWriteBuffer{i}.getMappedRange()).set("
              + f"_{inp_name});"
              + f"\n        gpuWriteBuffer{i}.unmap();\n        commandEncoder.copyBufferToBuffer(gpuWriteBuffer{i}, 0, {inp_name}, 0, gpuWriteBuffer{i}.size);"
              for i, inp_name in enumerate(input_names)
          ]
      )
      gpu_read_bufs = "\n    ".join(
          [
              f"const gpuReadBuffer{i} = device.createBuffer({{size:{output_name}.size, usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ }});"
              for i, output_name in enumerate(output_names)
          ]
      )
      outbuf_copies = "\n        ".join(
          [
              f"commandEncoder.copyBufferToBuffer({output_name}, 0, gpuReadBuffer{i}, 0, output{i}.size);"
              for i, output_name in enumerate(output_names)
          ]
      )
      output_readers = "\n        ".join(
          [
              f"await gpuReadBuffer{i}.mapAsync(GPUMapMode.READ);\n        const resultBuffer{i} = new Float32Array(gpuReadBuffer{i}.size);\n        resultBuffer{i}.set(new Float32Array(gpuReadBuffer{i}.getMappedRange()));\n        gpuReadBuffer{i}.unmap();"
              for i in range(len(output_names))
          ]
      )
      output_return = "[{}]".format(
          ",".join(
              [f"resultBuffer{i}" for i in range(len(output_names))]
          )
      )
      return (
          f"""
  {web_utils["getTensorBuffer"]}
  
  {web_utils["getTensorMetadata"]}
  
  const createEmptyBuf = (device, size) => {{
      return device.createBuffer({{size, usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST }});
  }};
  
  const createWeightBuf = (device, size, data) => {{
    const buf = device.createBuffer({{ mappedAtCreation: true, size, usage: GPUBufferUsage.STORAGE }});
    new Uint8Array(buf.getMappedRange()).set(data);
    buf.unmap();
    return buf;
  }};
  
  const addComputePass = (device, commandEncoder, pipeline, bufs, workgroup) => {{
    const bindGroup = device.createBindGroup({{layout: pipeline.getBindGroupLayout(0), entries: bufs.map((buffer, index) => ({{ binding: index, resource: {{ buffer }} }}))}});
    const passEncoder = commandEncoder.beginComputePass();
    passEncoder.setPipeline(pipeline);
    passEncoder.setBindGroup(0, bindGroup);
    passEncoder.dispatchWorkgroups(...workgroup);
    passEncoder.end();
  }};
  
  {kernel_code}
  
  const setupNet = async (device, safetensor) => {{
      const metadata = getTensorMetadata(safetensor);
  
      {_bufs}
  
      {gpu_write_bufs}
  
      {gpu_read_bufs}
  
      const kernels = [{kernel_names}];
      const piplines = await Promise.all(kernels.map(name => device.createComputePipelineAsync({{layout: "auto", compute: {{ module: device.createShaderModule({{ code: name }}), entryPoint: "main" }}}})));
  
      return async ({",".join([f"_{input_name}" for input_name in input_names])}) => {{
          const commandEncoder = device.createCommandEncoder();
          {input_writers}
          {kernel_calls}
          {outbuf_copies}
          const gpuCommands = commandEncoder.finish();
          device.queue.submit([gpuCommands]);
  
          {output_readers}
          return {output_return};
      }}
  }}
    """
          + f"\n\nconst loadNet = async (device) => {{ return await fetch('net.safetensors').then(x => x.arrayBuffer()).then(x => setupNet(device, new Uint8Array(x))); }}"
      )
  
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 5 for calling a global function with example usage: export_model_webgpu(functions, statements, bufs, bufs_to_save, weight_names, input_names, output_names) and returns: unknown