text: #! /usr/bin/env python
  # -*- coding: utf-8 -*-
  # vim:fenc=utf-8
  #
  # Distributed under terms of the MIT license.
  
  """An example of using tfp.optimizer.lbfgs_minimize to optimize a TensorFlow model.
  
  This code shows a naive way to wrap a tf.keras.Model and optimize it with the L-BFGS
  optimizer from TensorFlow Probability.
  
  Python interpreter version: 3.6.9
  TensorFlow version: 2.0.0
  TensorFlow Probability version: 0.8.0
  NumPy version: 1.17.2
  Matplotlib version: 3.1.1
  """
  import numpy
  import tensorflow as tf
  import tensorflow_probability as tfp
  from matplotlib import pyplot
  from . import complex_math 
  
  __all__ = ['function_factory']
  
  def function_factory(model, loss, dataset):
      """A factory to create a function required by tfp.optimizer.lbfgs_minimize.
  
      Args:
          model [in]: an instance of `tf.keras.Model` or its subclasses.
          loss [in]: a function with signature loss_value = loss(pred_y, true_y).
          train_x [in]: the input part of training data.
          train_y [in]: the output part of training data.
  
      Returns:
          A function that has a signature of:
              loss_value, gradients = f(model_parameters).
      """
  
      # obtain the shapes of all trainable parameters in the model
      shapes = tf.shape_n(model.trainable_variables)
      n_tensors = len(shapes)
  
      # we'll use tf.dynamic_stitch and tf.dynamic_partition later, so we need to
      # prepare required information first
      count = 0
      idx = [] # stitch indices
      part = [] # partition indices
  
      for i, shape in enumerate(shapes):
          n = numpy.product(shape)
          idx.append(tf.reshape(tf.range(count, count+n, dtype=tf.int32), shape))
          part.extend([i]*n)
          count += n
  
      part = tf.constant(part)
  
      @tf.function
      @tf.autograph.experimental.do_not_convert
      def assign_new_model_parameters(params_1d):
          """A function updating the model's parameters with a 1D tf.Tensor.
  
          Args:
              params_1d [in]: a 1D tf.Tensor representing the model's trainable parameters.
          """
  
          params = tf.dynamic_partition(params_1d, part, n_tensors)
          for i, (shape, param) in enumerate(zip(shapes, params)):
              model.trainable_variables[i].assign(tf.reshape(param, shape))
              #tf.print(model.trainable_variables[i])
  
      @tf.function
      def volume_form(x, Omega_Omegabar, mass, restriction):
  
          kahler_metric = complex_math.complex_hessian(tf.math.real(model(x)), x)
          volume_form = tf.math.real(tf.linalg.det(tf.matmul(restriction, tf.matmul(kahler_metric, restriction, adjoint_b=True))))
          weights = mass / tf.reduce_sum(mass)
          factor = tf.reduce_sum(weights * volume_form / Omega_Omegabar)
          #factor = tf.constant(35.1774, dtype=tf.complex64)
          return volume_form / factor
  
  
      # now create a function that will be returned by this factory
      def f(params_1d):
          """A function that can be used by tfp.optimizer.lbfgs_minimize.
  
          This function is created by function_factory.
  
          Args:
             params_1d [in]: a 1D tf.Tensor.
  
          Returns:
              A scalar loss and the gradients w.r.t. the `params_1d`.
          """
  
          # use GradientTape so that we can calculate the gradient of loss w.r.t. parameters
          for step, (points, Omega_Omegabar, mass, restriction) in enumerate(dataset):
              with tf.GradientTape() as tape:
                  # update the parameters in the model
                  assign_new_model_parameters(params_1d)
                  # calculate the loss
                  det_omega = volume_form(points, Omega_Omegabar, mass, restriction)
                  loss_value = loss(Omega_Omegabar, det_omega, mass)
  
              # calculate gradients and convert to 1D tf.Tensor
              grads = tape.gradient(loss_value, model.trainable_variables)
              grads = tf.dynamic_stitch(idx, grads)
            
              # reweight the loss and grads 
              mass_sum = tf.reduce_sum(mass) 
              try:
                  total_loss += loss_value * mass_sum 
                  total_grads += grads * mass_sum
                  total_mass += mass_sum
              except NameError:
                  total_loss = loss_value * mass_sum 
                  total_grads = grads * mass_sum
                  total_mass = mass_sum
  
          total_loss = total_loss / total_mass
          total_grads = total_grads / total_mass
  
          # print out iteration & loss
          f.iter.assign_add(1)
          tf.print("Iter:", f.iter, "loss:", total_loss)
  
          # store loss value so we can retrieve later
          tf.py_function(f.history.append, inp=[total_loss], Tout=[])
  
          return total_loss, total_grads
  
      # store these information as members so we can use them outside the scope
      f.iter = tf.Variable(0)
      f.idx = idx
      f.part = part
      f.shapes = shapes
      f.assign_new_model_parameters = assign_new_model_parameters
      f.history = []
  
      return f
  