code: def transport_sample(
      all_preds,
      label,
      choice,
      prob,
      hist_loss,
      count,
      transport_method,
      alpha,
      training=False,
  ):
      """
      sample-wise transport
  
      Args:
          all_preds (torch.Tensor): predictions from all predictors, [sample x states]
          label (torch.Tensor): label, [sample]
          choice (torch.Tensor): gumbel softmax choice, [sample x states]
          prob (torch.Tensor): router predicted probility, [sample x states]
          hist_loss (torch.Tensor): history loss matrix, [sample x states]
          count (list): sample counts for each day, empty list for sample-wise transport
          transport_method (str): transportation method
          alpha (float): fusion parameter for calculating transport loss matrix
          training (bool): indicate training or inference
      """
      assert all_preds.shape == choice.shape
      assert len(all_preds) == len(label)
      assert transport_method in ["oracle", "router"]
  
      all_loss = torch.zeros_like(all_preds)
      mask = ~torch.isnan(label)
      all_loss[mask] = (all_preds[mask] - label[mask, None]).pow(
          2
      )  # [sample x states]
  
      L = minmax_norm(all_loss.detach())
      Lh = L * alpha + minmax_norm(hist_loss) * (
          1 - alpha
      )  # add hist loss for transport
      Lh = minmax_norm(Lh)
      P = sinkhorn(-Lh)
      del Lh
  
      if transport_method == "router":
          if training:
              pred = (all_preds * choice).sum(dim=1)  # gumbel softmax
          else:
              pred = all_preds[
                  range(len(all_preds)), prob.argmax(dim=-1)
              ]  # argmax
      else:
          pred = (all_preds * P).sum(dim=1)
  
      if transport_method == "router":
          loss = loss_fn(pred, label)
      else:
          loss = (all_loss * P).sum(dim=1).mean()
  
      return loss, pred, L, P
  
imports: import io
  import os
  import copy
  import math
  import json
  import numpy
  import pandas
  import matplotlib.pyplot
  import torch
  import torch.nn
  import torch.optim
  import torch.nn.functional
  from tqdm import tqdm
  from qlib.constant import EPS
  from qlib.log import get_module_logger
  from qlib.model.base import Model
  from qlib.contrib.data.dataset import MTSDatasetH
label_desc: Use image node_id 6 for calling a global function with example usage: transport_sample(all_preds, label, choice, prob, hist_loss, count, transport_method, alpha, training) and returns: loss,  pred,  L,  P