// Shape inference should not change the element type.
      if (HasCompatibleElementTypes(input.getType(), result.getType()) &&
          meet == input.getType()) {
        arg_op.set(input);
      } else {
        OpBuilder b(return_op.getOperation());
        auto new_cast_op = InsertCast(b, return_op.getLoc(), meet, input);
        if (!new_cast_op) return failure();
        arg_op.set(new_cast_op->getResult(0));

    module {
      func @while_body {
        %input = "tf.IteratorGetNext"(...) {device = "/CPU:0"}: -> tensor<2x224x224x3xf32>
        %device_launch = "tf_device.cluster_func"(%input,...) {func = @_func,...)
        return ...
      }
      func @_func(%input: tensor<2x224x224x3xf32>, %filter: tensor<7x7x3x64xf32>) {

          dtypes.float32, shape=[], name='input'
      )
      q_input = array_ops.fake_quant_with_min_max_args(
          input_tensor, min=-0.1, max=0.2, num_bits=8, narrow_range=False
      )
      sqrt = math_ops.sqrt(q_input, name='sqrt')
      identity = array_ops.identity(sqrt, name='output')

      input_map = {'input': input_tensor}
      output_map = {'sqrt': identity}

void TF_AddInput(TF_OperationDescription* desc, TF_Output input) {
  desc->node_builder.Input(&input.oper->node, input.index);
}

void TF_AddInputList(TF_OperationDescription* desc, const TF_Output* inputs,
                     int num_inputs) {
  std::vector<NodeBuilder::NodeOut> input_list;
  input_list.reserve(num_inputs);
  for (int i = 0; i < num_inputs; ++i) {
    input_list.emplace_back(&inputs[i].oper->node, inputs[i].index);
  }

  // expected-remark@above {{ID: 7}}
  %input: tensor<f32>,
  %group_key: tensor<i32>,
  %group_size: tensor<i32>,
  %instance_key: tensor<i32>) {
  tf_executor.graph {
    // expected-remark@above {{ID: 5}}
    %island = tf_executor.island {
        // expected-remark@above {{ID: 3}}
        // expected-remark@above {{Successors: {4}}}

  Value reshaped_input = builder.create<TF::ReshapeOp>(
      RankedTensorType::get(shape, element_type), input, shape_cst);
  return reshaped_input;
}

// Create a tensor which is equal to input[begin: begin + size].
Value BuildSliceOp(ImplicitLocOpBuilder& builder,
                   ConversionPatternRewriter& rewriter, Value input,
                   Value begin, ArrayRef<int64_t> shape, Type idx_type,

    func.return %1, %3 : tensor<*xi32>, tensor<*xi1>
  }
}

// -----

// Tests inputs are correctly split and fed into TPU computation for tiled input
// sharding.

// The following OpSharding is used for TPU computation inputs in below test:
// Proto debug string:
//  input 0
//   type: OTHER
//   tile_assignment_dimensions: 1
//   tile_assignment_dimensions: 2

// RUN: tf-opt -split-input-file -verify-diagnostics -tfl-runtime-verify %s | FileCheck %s

// Unary math ops
// -----

// CHECK-LABEL: testCos
func.func @testCos(tensor<? x f32>) -> tensor<? x f32> {
^bb0(%arg0: tensor<? x f32>):
  // CHECK: "tfl.cos"(%arg0)
  %0 = "tfl.cos"(%arg0): (tensor<? x f32>) -> tensor<? x f32>
  func.return %0 : tensor<? x f32>
}

// -----

// test invalid Cos input

    func.return %1 : tensor<2xi32>
  }

  //  Tests the case when an outside compiled op with some dynamically shaped input but static output
  //  comes before an outside compiled op with some dynamically shaped input/output.  This ensures that
  //  for the second outside compiled op, all of the inputs are explicitly provided.

  // CHECK-LABEL: func @static_output_dynamic_input_outside_compilation

    protected A doBackward(B b) {
      return convert(bimap.inverse(), b);
    }

    private static <X, Y> Y convert(BiMap<X, Y> bimap, X input) {
      Y output = bimap.get(input);
      checkArgument(output != null, "No non-null mapping present for input: %s", input);
      return output;
    }

    @Override
    public boolean equals(@CheckForNull Object object) {