code: def test_custom_trace_fn(self):
      init_x = np.array([0.0, 0.0]).astype(np.float32)
      target_x = np.array([3.0, 4.0]).astype(np.float32)
      loss_fn = lambda x: tf.reduce_sum((x - target_x) ** 2)
  
      # The trace_fn should determine the structure and values of the results.
      def trace_fn(traceable_quantities):
          return {
              "loss": traceable_quantities.loss,
              "x": traceable_quantities.parameters,
              "sqdiff": (traceable_quantities.parameters - target_x)
              ** 2,
          }
  
      _, results = minimize_stateless(
          loss_fn,
          init=init_x,
          num_steps=100,
          optimizer=StatelessSGD(0.05),
          trace_fn=trace_fn,
      )
      results_ = self.evaluate(results)
      self.assertAllClose(results_["x"][0], init_x, atol=0.5)
      self.assertAllClose(results_["x"][-1], target_x, atol=0.2)
      self.assertAllClose(results_["sqdiff"][-1], [0.0, 0.0], atol=0.1)
  
imports: import collections
  import numpy
  import tensorflow.compat.v1
  import tensorflow.compat.v2
  from tensorflow_probability.python import optimizer
  from tensorflow_probability.python.distributions import normal
  from tensorflow_probability.python.internal import test_util
  from tensorflow_probability.python.internal import tf_keras
  from tensorflow_probability.python.math.minimize import minimize
  from tensorflow_probability.python.math.minimize import minimize_stateless
label_desc: Use image node_id 6 for calling the MinimizeStatelessTests obj's underlying member method code with example usage: obj.test_custom_trace_fn() and returns: dict