class_docstr: Testing rig.
class_docstr_tok: ["Testing","rig","."]
docstr: Tests that the MST op can recover a simple tree.
docstr_tok: ["Tests","that","the","MST","op","can","recover","a","simple","tree","."]
code: def testMaximumSpanningTree(self):
      """Tests that the MST op can recover a simple tree."""
      # The first batch element prefers 3 as root, then 3->0->1->2, for a total
      # score of 4+2+1=7.  The second batch element is smaller and has reversed
      # scores, so 0 is root and 0->2->1.
      num_nodes = constant_op.constant([4, 3], dtypes.int32)
      scores = constant_op.constant(
          [
              [[0, 0, 0, 0], [1, 0, 0, 0], [1, 2, 0, 0], [1, 2, 3, 4]],
              [[4, 3, 2, 9], [0, 0, 2, 9], [0, 0, 0, 9], [9, 9, 9, 9]],
          ],
          dtypes.int32,
      )  # pyformat: disable
  
      (max_scores, argmax_sources) = mst_ops.max_spanning_tree(
          num_nodes, scores, forest=False
      )
  
      self.assertAllEqual(max_scores, [7, 6])
      self.assertAllEqual(argmax_sources, [[3, 0, 1, 3], [0, 2, 0, -1]])
  
file_path: text/tensorflow_text/python/ops/mst_ops_test.py
imports: import numpy
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import test_util
  from tensorflow.python.platform import test
  from tensorflow_text.python.ops import mst_ops
label_desc: Use image node_id 1 for calling the MstOpsTest obj's underlying member method code with example usage: obj.testMaximumSpanningTree() without return types