code: def test_multiplex_2d(self):
      a = tf.constant([[1, 2, 3], [4, 5, 6]], dtype=tf.int64)
      b = tf.constant([[10, 20, 30], [40, 50, 60]], dtype=tf.int64)
      cond = tf.constant(
          [[True, False, True], [False, True, False]], dtype=bool
      )
      expect = np.where(
          self.evaluate(cond), self.evaluate(a), self.evaluate(b)
      )
      # expected result is [[1, 20], [3, 40]]
      result = multiplex_2_op.multiplex(cond, a, b)
      self.assertAllEqual(result, expect)
  
file_path: tensorflow-upstream/tensorflow/examples/custom_ops_doc/multiplex_2/multiplex_2_test.py
imports: import numpy
  import tensorflow
  from tensorflow.examples.custom_ops_doc.multiplex_2 import multiplex_2_op
  from tensorflow.python.framework import errors_impl
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 5 for calling the MultiplexOpRank1Test obj's underlying member method code with example usage: obj.test_multiplex_2d() without return types