message: Fix failing xla test in v2
  Summary:
  Ref T16509
  TF2.1 Only
  Test Plan: It passes
  Reviewers: hakons, alfiee, #tensorflow, #framework_ip_review_-_any_oss_or_third-party_code_use_has_been_approved
  Maniphest Tasks: T16509
diff: @@ -400,6 +400,7 @@ class XlaOpsShapeInferenceTest(xla_test.XLATestCase, parameterized.TestCase):
  xla.dot_general(a, b, dim_nums)
  def testDotShapeInference(self):
  + with self.session():
  a = array_ops.placeholder(np.float32, shape=(1, 2, 3, 4))
  b = array_ops.placeholder(np.float32, shape=(4, 3, 2, 1))
  