code: def umap_loss(
      batch_size,
      negative_sample_rate,
      _a,
      _b,
      edge_weights,
      parametric_embedding,
      repulsion_strength=1.0,
  ):
      """
      Generate a keras-compatible loss function for UMAP loss
  
      Parameters
      ----------
      batch_size : int
          size of mini-batches
      negative_sample_rate : int
          number of negative samples per positive samples to train on
      _a : float
          distance parameter in embedding space
      _b : float
          distance parameter in embedding space
      edge_weights : array
          weights of all edges from sparse UMAP graph
      parametric_embedding : bool
          whether the embedding is parametric or nonparametric
      repulsion_strength : float, optional
          strength of repulsion vs attraction for cross-entropy, by default 1.0
  
      Returns
      -------
      loss : function
          loss function that takes in a placeholder (0) and the output of the keras network
      """
  
      if not parametric_embedding:
          # multiply loss by weights for nonparametric
          weights_tiled = np.tile(
              edge_weights, negative_sample_rate + 1
          )
  
      @tf.function
      def loss(placeholder_y, embed_to_from):
          # split out to/from
          embedding_to, embedding_from = tf.split(
              embed_to_from, num_or_size_splits=2, axis=1
          )
  
          # get negative samples
          embedding_neg_to = tf.repeat(
              embedding_to, negative_sample_rate, axis=0
          )
          repeat_neg = tf.repeat(
              embedding_from, negative_sample_rate, axis=0
          )
          embedding_neg_from = tf.gather(
              repeat_neg,
              tf.random.shuffle(tf.range(tf.shape(repeat_neg)[0])),
          )
  
          #  distances between samples (and negative samples)
          distance_embedding = tf.concat(
              [
                  tf.norm(embedding_to - embedding_from, axis=1),
                  tf.norm(
                      embedding_neg_to - embedding_neg_from, axis=1
                  ),
              ],
              axis=0,
          )
  
          # convert distances to probabilities
          log_probabilities_distance = (
              convert_distance_to_log_probability(
                  distance_embedding, _a, _b
              )
          )
  
          # set true probabilities based on negative sampling
          probabilities_graph = tf.concat(
              [
                  tf.ones(batch_size),
                  tf.zeros(batch_size * negative_sample_rate),
              ],
              axis=0,
          )
  
          # compute cross entropy
          (
              attraction_loss,
              repellant_loss,
              ce_loss,
          ) = compute_cross_entropy(
              probabilities_graph,
              log_probabilities_distance,
              repulsion_strength=repulsion_strength,
          )
  
          if not parametric_embedding:
              ce_loss = ce_loss * weights_tiled
  
          return tf.reduce_mean(ce_loss)
  
      return loss
  
imports: import numpy
  from umap import UMAP
  from warnings import warn, catch_warnings, filterwarnings
  from numba import TypingError
  import os
  from umap.spectral import spectral_layout
  from sklearn.utils import check_random_state
  import codecs, pickle
  from sklearn.neighbors import KDTree
  import sys
label_desc: Use image node_id 5 for calling a global function with example usage: umap_loss(batch_size, negative_sample_rate, _a, _b, edge_weights, parametric_embedding, repulsion_strength) and returns: loss, tf