code: def test_transform_fn_train(self, loss_key):
      estimator_with_default_transform_fn = (
          self._create_default_estimator(loss_key)
      )
  
      # The below tests the `transform_fn` in the TRAIN mode. In this mode, the
      # `_transform_fn` invokes the `encode_listwise_features()`, which requires
      # 3D example features and 2D context features.
      (
          context,
          example,
      ) = estimator_with_default_transform_fn._transform_fn(
          {
              "f1": tf.ones([10, 10, 1], dtype=tf.float32),
              "f2": tf.ones([10, 10, 1], dtype=tf.float32) * 2.0,
              "f3": tf.ones([10, 10, 1], dtype=tf.float32) * 3.0,
              "c1": tf.ones([10, 1], dtype=tf.float32),
              "c2": tf.ones([10, 1], dtype=tf.float32) * 2.0,
          },
          tf_estimator.ModeKeys.TRAIN,
      )
      # `c1` is the only context feature defined in `_context_feature_columns()`.
      self.assertCountEqual(context.keys(), ["c1"])
  
      # `f1`, `f2`, `f3` are all defined in the `_example_feature_columns()`.
      self.assertCountEqual(example.keys(), ["f1", "f2", "f3"])
  
      # Validates the `context` and `example` features are transformed correctly.
      self.assertAllEqual(tf.ones(shape=[10, 1]), context["c1"])
      self.assertAllEqual(tf.ones(shape=[10, 10, 1]), example["f1"])
  
imports: import os
  from absl.testing import parameterized
  import six
  import tensorflow
  from tensorflow import estimator
  from google.protobuf import text_format
  from tensorflow_ranking.python import data
  from tensorflow_ranking.python import estimator
  from tensorflow_ranking.python import feature
  from tensorflow_serving.apis import input_pb2
label_desc: Use image node_id 4 for calling the EstimatorBuilderTest obj's underlying member method code with example usage: obj.test_transform_fn_train(loss_key) without return types