// XLA_PerChannel-DAG: %[[out_scale:.*]] = "tf.Const"() <{value = dense<7.24974147E-4> : tensor<f32>}> : () -> tensor<f32>
}

// TODO(b/308773062): Add whole_model unit-test

// -----

module {
  func.func @matmul2_with_int_per_layer(%arg0: tensor<2x2xf32> {tf_saved_model.index_path = ["input_tensor"]}) -> tensor<2x2xf32> {

          outputs={output_key: output_tensor},
      )

    return in_placeholder

  def _create_while_model(self, input_shape: Sequence[int] = (1, 32, 32, 512)):
    class WhileModel(module.Module):
      """A model with a while op."""

      def __init__(self):
        w_shape = [3, 3] + [input_shape[-1], input_shape[-1]]