code: def loss(
      example_embedding, context_embedding, negative_embedding, manifold
  ):
      """Compute loss and grad.
  
      Compute loss and grad given embedding of the current example,
      embedding of the context and negative sampling embedding.
      """
      n_edges, dim = (
          negative_embedding.shape[0],
          example_embedding.shape[-1],
      )
      example_embedding = gs.expand_dims(example_embedding, 0)
      context_embedding = gs.expand_dims(context_embedding, 0)
  
      positive_distance = manifold.metric.squared_dist(
          example_embedding, context_embedding
      )
      positive_loss = log_sigmoid(-positive_distance)
  
      reshaped_example_embedding = gs.repeat(
          example_embedding, n_edges, axis=0
      )
  
      negative_distance = manifold.metric.squared_dist(
          reshaped_example_embedding, negative_embedding
      )
      negative_loss = log_sigmoid(negative_distance)
  
      total_loss = -(positive_loss + negative_loss.sum())
  
      positive_log_sigmoid_grad = -grad_log_sigmoid(-positive_distance)
  
      positive_distance_grad = grad_squared_distance(
          example_embedding, context_embedding
      )
  
      positive_grad = (
          gs.repeat(positive_log_sigmoid_grad, dim, axis=-1)
          * positive_distance_grad
      )
  
      negative_distance_grad = grad_squared_distance(
          reshaped_example_embedding, negative_embedding
      )
  
      negative_distance = gs.to_ndarray(
          negative_distance, to_ndim=2, axis=-1
      )
      negative_log_sigmoid_grad = grad_log_sigmoid(negative_distance)
  
      negative_grad = negative_log_sigmoid_grad * negative_distance_grad
  
      example_grad = -(positive_grad + negative_grad.sum(axis=0))
  
      return total_loss, example_grad
  
file_path: geomstats/examples/learning_graph_structured_data_h2.py
imports: import logging
  import matplotlib.pyplot
  import geomstats.backend
  import geomstats.visualization
  from geomstats.datasets.utils import load_karate_graph
  from geomstats.geometry.poincare_ball import PoincareBall
label_desc: Use image node_id 4 for calling a global function with example usage: loss(example_embedding, context_embedding, negative_embedding, manifold) and returns: total_loss,  example_grad