code: def optimal_transport_dist(
      txt_emb: Tensor,
      img_emb: Tensor,
      txt_pad: Tensor,
      img_pad: Tensor,
      beta: float = 0.5,
      iteration: int = 50,
      k: int = 1,
  ) -> Tensor:
      """[B, M, D], [B, N, D], [B, M], [B, N]"""
      cost = cost_matrix_cosine(txt_emb, img_emb)
      # mask the padded inputs
      joint_pad = txt_pad.unsqueeze(-1) | img_pad.unsqueeze(-2)
      cost.masked_fill_(joint_pad, 0)
  
      txt_len = (
          txt_pad.size(1) - txt_pad.sum(dim=1, keepdim=False)
      ).to(dtype=cost.dtype)
      img_len = (
          img_pad.size(1) - img_pad.sum(dim=1, keepdim=False)
      ).to(dtype=cost.dtype)
  
      T = ipot(
          cost.detach(),
          txt_len,
          txt_pad,
          img_len,
          img_pad,
          joint_pad,
          beta,
          iteration,
          k,
      )
      distance = trace(cost.matmul(T.detach()))
      return distance
  
imports: import torch
  from torch import Tensor
  from torch.nn import functional
label_desc: Use image node_id 4 for calling a global function with example usage: optimal_transport_dist(txt_emb, img_emb, txt_pad, img_pad, beta, iteration, k) and returns: distance