code: def train_model():
      """
      Train the patch similarity function
      """
      global ema, model
  
      model = Model()
  
      def loss(x, y):
          """
          K-way contrastive loss as in SimCLR et al.
          The idea is that we should embed x and y so that they are similar
          to each other, and dis-similar from others. To do this we have a
          softmx loss over one dimension to make the values large on the diagonal
          and small off-diagonal.
          """
          a = model.encode(x)
          b = model.decode(y)
  
          mat = a @ b.T
          return objax.functional.loss.cross_entropy_logits_sparse(
              logits=jn.exp(jn.clip(model.scale.w.value, -2, 4)) * mat,
              labels=np.arange(a.shape[0]),
          ).mean()
  
      ema = objax.optimizer.ExponentialMovingAverage(
          model.vars(), momentum=0.999
      )
      gv = objax.GradValues(loss, model.vars())
  
      encode_ema = ema.replace_vars(lambda x: model.encode(x))
      decode_ema = ema.replace_vars(lambda y: model.decode(y))
  
      def train_op(x, y):
          """
          No one was ever fired for using Adam with 1e-4.
          """
          g, v = gv(x, y)
          opt(1e-4, g)
          ema()
          return v
  
      opt = objax.optimizer.Adam(model.vars())
      train_op = objax.Jit(
          train_op, gv.vars() + opt.vars() + ema.vars()
      )
  
      ys_ = ys_train
  
      print(ys_.shape)
  
      xs_ = xs_train.reshape((-1, xs_train.shape[-1]))
      ys_ = ys_.reshape((-1, ys_train.shape[-1]))
  
      # The model scale trick here is taken from CLIP.
      # Let the model decide how confident to make its own predictions.
      model.scale.w.assign(jn.zeros((1, 1)))
  
      valid_size = 1000
  
      print(xs_train.shape)
      # SimCLR likes big batches
      B = 4096
      for it in range(80):
          print()
          ms = []
          for i in range(1000):
              # First batch is smaller, to make training more stable
              bs = [B // 64, B][it > 0]
              batch = np.random.randint(
                  0, len(xs_) - valid_size, size=bs
              )
              r = train_op(xs_[batch], ys_[batch])
  
              # This shouldn't happen, but if it does, better to bort early
              if np.isnan(r):
                  print("Die on nan")
                  print(ms[-100:])
                  return
              ms.append(r)
  
          print("mean", np.mean(ms), "scale", model.scale.w.value)
          print("loss", loss(xs_[-100:], ys_[-100:]))
  
          a = encode_ema(xs_[-valid_size:])
          b = decode_ema(ys_[-valid_size:])
  
          br = b[np.random.permutation(len(b))]
  
          print(
              "score",
              np.mean(
                  np.sum(a * b, axis=(1)) - np.sum(a * br, axis=(1))
              ),
              np.mean(
                  np.sum(a * b, axis=(1)) > np.sum(a * br, axis=(1))
              ),
          )
      ckpt = objax.io.Checkpoint("saved", keep_ckpts=0)
      ema.replace_vars(lambda: ckpt.save(model.vars(), 0))()
  
imports: import collections
  import hashlib
  import time
  import multiprocessing
  import torch
  import numpy
  import torch.nn
  import scipy.stats
  import matplotlib.pyplot
  from PIL import Image
  import jax
  import jax.numpy
  import objax
  import scipy.optimize
  import numpy
  import multiprocessing
label_desc: Use image node_id 6 for calling a global function with example usage: train_model() and returns: objax, v