code: def test_get_train_op(self):
      with tf.Graph().as_default():
          logits = tf.cast([[1.0]], tf.float32)
          labels = tf.cast([[0.0]], tf.float32)
          loss = tf.abs(tf.reduce_sum(input_tensor=logits - labels))
  
          def _train_op_fn(loss):
              tf.debugging.assert_near(
                  tf.cast(loss, dtype=tf.float32),
                  1.0,
                  name="assert_loss",
              )
              return tf.constant(b"train_op_fn")
  
          class _Optimizer(object):
              def minimize(self, loss, global_step):
                  del global_step
                  tf.debugging.assert_near(
                      tf.cast(loss, dtype=tf.float32),
                      1.0,
                      name="assert_loss",
                  )
                  return tf.constant(b"optimizer")
  
          train_op = ranking_head._get_train_op(
              loss, _train_op_fn, None
          )
          self.assertIsNotNone(train_op)
          with self.cached_session() as sess:
              train_result = sess.run(train_op)
              self.assertEqual(b"train_op_fn", train_result)
  
          train_op = ranking_head._get_train_op(
              loss, None, _Optimizer()
          )
          self.assertIsNotNone(train_op)
          with self.cached_session() as sess:
              train_result = sess.run(train_op)
              self.assertEqual(b"optimizer", train_result)
  
          with self.assertRaisesRegex(
              ValueError,
              r"train_op_fn and optimizer cannot both be set.",
          ):
              ranking_head._get_train_op(
                  loss, _train_op_fn, _Optimizer()
              )
          with self.assertRaisesRegex(
              ValueError,
              r"train_op_fn and optimizer cannot both be None.",
          ):
              ranking_head._get_train_op(loss, None, None)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import tensorflow
  from tensorflow import estimator
  from tensorflow_ranking.python import head
  from tensorflow_ranking.python import metrics
label_desc: Use image node_id 1 for calling the UtilTest obj's underlying member method code with example usage: obj.test_get_train_op() and returns: tf, tf