}

func.func @RemoveTrivialAddBf16RHS(%arg0: tensor<2x2xbf16>) -> tensor<2x2xbf16> {
  %cst = arith.constant dense<0.0> : tensor<2x2xbf16>
  %0 = "tf.Add"(%arg0, %cst) : (tensor<2x2xbf16>, tensor<2x2xbf16>) -> tensor<2x2xbf16>
  func.return %0 : tensor<2x2xbf16>

  // CHECK-LABEL: RemoveTrivialAdd
  // CHECK-NEXT: return %arg0 : tensor<2x2xbf16>
}

  func.return %0: tensor<3x2xf16>
}

func.func @sparse_qu8() -> tensor<3x2x!quant.uniform<u8<1:255>:f32, 1.0>> {
  // CHECK-LABEL: @sparse_qu8

      return %0 : tensor<1x?x4xbf16>
    }
  })";

  TF_ASSERT_OK(LegalizeModule(kModuleWithNonConstParam));
}

TEST_F(Tf2XlaRewriterTest, ErrorsWithInvalidNumberOfParametersToArgs) {