code: def construct_edge_dataset(
      X,
      graph_,
      n_epochs,
      batch_size,
      parametric_embedding,
      parametric_reconstruction,
      global_correlation_loss_weight,
  ):
      """
      Construct a tf.data.Dataset of edges, sampled by edge weight.
  
      Parameters
      ----------
      X : array, shape (n_samples, n_features)
          New data to be transformed.
      graph_ : scipy.sparse.csr.csr_matrix
          Generated UMAP graph
      n_epochs : int
          # of epochs to train each edge
      batch_size : int
          batch size
      parametric_embedding : bool
          Whether the embedder is parametric or non-parametric
      parametric_reconstruction : bool
          Whether the decoder is parametric or non-parametric
      """
  
      def gather_index(index):
          return X[index]
  
      # if X is > 512Mb in size, we need to use a different, slower method for
      #    batching data.
      gather_indices_in_python = (
          True if X.nbytes * 1e-9 > 0.5 else False
      )
  
      def gather_X(edge_to, edge_from):
          # gather data from indexes (edges) in either numpy of tf, depending on array size
          if gather_indices_in_python:
              edge_to_batch = tf.py_function(
                  gather_index, [edge_to], [tf.float32]
              )[0]
              edge_from_batch = tf.py_function(
                  gather_index, [edge_from], [tf.float32]
              )[0]
          else:
              edge_to_batch = tf.gather(X, edge_to)
              edge_from_batch = tf.gather(X, edge_from)
          return edge_to_batch, edge_from_batch
  
      def get_outputs(edge_to_batch, edge_from_batch):
          outputs = {"umap": tf.repeat(0, batch_size)}
          if global_correlation_loss_weight > 0:
              outputs["global_correlation"] = edge_to_batch
          if parametric_reconstruction:
              # add reconstruction to iterator output
              # edge_out = tf.concat([edge_to_batch, edge_from_batch], axis=0)
              outputs["reconstruction"] = edge_to_batch
          return (edge_to_batch, edge_from_batch), outputs
  
      def make_sham_generator():
          """
          The sham generator is a placeholder when all data is already intrinsic to
          the model, but keras wants some input data. Used for non-parametric
          embedding.
          """
  
          def sham_generator():
              while True:
                  yield tf.zeros(1, dtype=tf.int32), tf.zeros(
                      1, dtype=tf.int32
                  )
  
          return sham_generator
  
      # get data from graph
      (
          graph,
          epochs_per_sample,
          head,
          tail,
          weight,
          n_vertices,
      ) = get_graph_elements(graph_, n_epochs)
  
      # number of elements per batch for embedding
      if batch_size is None:
          # batch size can be larger if its just over embeddings
          if parametric_embedding:
              batch_size = np.min([n_vertices, 1000])
          else:
              batch_size = len(head)
  
      edges_to_exp, edges_from_exp = (
          np.repeat(head, epochs_per_sample.astype("int")),
          np.repeat(tail, epochs_per_sample.astype("int")),
      )
  
      # shuffle edges
      shuffle_mask = np.random.permutation(range(len(edges_to_exp)))
      edges_to_exp = edges_to_exp[shuffle_mask].astype(np.int64)
      edges_from_exp = edges_from_exp[shuffle_mask].astype(np.int64)
  
      # create edge iterator
      if parametric_embedding:
          edge_dataset = tf.data.Dataset.from_tensor_slices(
              (edges_to_exp, edges_from_exp)
          )
          edge_dataset = edge_dataset.repeat()
          edge_dataset = edge_dataset.shuffle(10000)
          edge_dataset = edge_dataset.batch(
              batch_size, drop_remainder=True
          )
          edge_dataset = edge_dataset.map(
              gather_X, num_parallel_calls=tf.data.experimental.AUTOTUNE
          )
          edge_dataset = edge_dataset.map(
              get_outputs,
              num_parallel_calls=tf.data.experimental.AUTOTUNE,
          )
          edge_dataset = edge_dataset.prefetch(10)
      else:
          # nonparametric embedding uses a sham dataset
          gen = make_sham_generator()
          edge_dataset = tf.data.Dataset.from_generator(
              gen,
              (tf.int32, tf.int32),
              output_shapes=(tf.TensorShape(1), tf.TensorShape((1,))),
          )
      return (
          edge_dataset,
          batch_size,
          len(edges_to_exp),
          head,
          tail,
          weight,
      )
  
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 8 for calling a global function with example usage: construct_edge_dataset(X, graph_, n_epochs, batch_size, parametric_embedding, parametric_reconstruction, global_correlation_loss_weight) and returns: edge_dataset,  batch_size,  len,  head,  tail,  weight, X, edge_to_batch,  edge_from_batch, edge_to_batch,  edge_from_batch,  outputs, sham_generator