%0 = "tf.D"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
    %1 = func.call @referenced_func2(%0) : (tensor<?xi32>) -> tensor<?xi32>
    func.return %1 : tensor<?xi32>
  }

  func.func @referenced_func2(%arg0: tensor<?xi32>) -> tensor<?xi32> {
    %0 = "tf.E"(%arg0) : (tensor<?xi32>) -> tensor<?xi32>
    func.return %0 : tensor<?xi32>
  }
}

// -----

    %4 = "tf_device.launch"() ({
      %c = "tf.opC"(%1) {device = "TPU_REPLICATED_CORE_0"} : (tensor<?xi32>) -> tensor<*xi32>
      tf_device.return %c : tensor<*xi32>
    }) {device = "c"} : () -> tensor<*xi32>
    tf_device.return %1, %2, %3, %4 : tensor<?xi32>, tensor<*xi32>, tensor<*xi32>, tensor<*xi32>
  }
  func.return
}

// CHECK:  [[SHAPE:%.*]] = "tf.Shape"([[VAL_0]])

        "tf.AssignVariableOp"(%1, %3) {_tpu_replicate = "cluster"} : (tensor<!tf_type.resource<tensor<i32>>>, tensor<*xi32>) -> ()
        %4:2 = "tf.TPUReplicatedOutput"(%1) : (tensor<!tf_type.resource<tensor<i32>>>) -> (tensor<!tf_type.resource<tensor<i32>>>, tensor<!tf_type.resource<tensor<i32>>>)
        tf_executor.yield %arg0 : tensor<!tf_type.resource<tensor<i32>>>
      }

      %1 = "tfl.greater"(%arg2, %cst_0) : (tensor<*xi32>, tensor<i32>) -> tensor<i1>
      "tfl.yield"(%1) : (tensor<i1>) -> ()
  },  {
    ^bb0(%arg2: tensor<*xi32>, %arg3: tensor<*xf32>, %arg4: tensor<i32>):
      %1 = "tfl.sub"(%arg2, %arg4) {fused_activation_function = "NONE"} :
        (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
      %2 = tfl.add %arg3, %arg3 {fused_activation_function = "NONE"} : tensor<*xf32>

}

func.func @"__inference_Dataset_map_<lambda>_80"(%arg0: tensor<*xi32>) -> tensor<*xi32> {
  %0 = "tf.Const"() {value = dense<2> : tensor<i32>} : () -> tensor<i32>
  %1 = "tf.Mul"(%arg0, %0) {device = ""} : (tensor<*xi32>, tensor<i32>) -> tensor<*xi32>
  %2 = "tf.Identity"(%1) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>
  func.return %2 : tensor<*xi32>
}

tensor<i32> loc(fused[#loc1, #loc2, #loc3]), %arg3: tensor<2xi32> loc(fused[#loc1, #loc2, #loc3])) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1120x?xi32>, tensor<2xi32>) {\0A    %0 = \22tf.Reshape\22(%arg0, %arg1) {_xla_outside_compilation = \220\22} : (tensor<3360x?xi32>, tensor<3xi32>) -> tensor<3x1120x?xi32> loc(#loc9)\0A    %1:3 = \22tf.Split\22(%arg2, %0) {_xla_outside_compilation = \220\22} : (tensor<i32>, tensor<3x1120x?xi32>) -> (tensor<1x1120x?xi32>, tensor<1x1120x?xi32>, tensor<1x1120x?xi32>)...

    // CHECK: tf.BroadcastGradientArgs
    // CHECK-SAME: (tensor<*xi32>, tensor<0xi32>) -> (tensor<?xi32>, tensor<?xi32>)
    %2:2 = "tf.BroadcastGradientArgs"(%arg0, %arg1) {T = "tfdtype$DT_INT32", name = "BroadcastGradientArgs"} : (tensor<*xi32>, tensor<0xi32>) -> (tensor<?xi32>, tensor<?xi32>)
    func.return %2#0 : tensor<?xi32>
  }

// CHECK:           return %[[VAL_4]] : tensor<?xi32>
// CHECK:         }
func.func @relu6_unranked(%arg0: tensor<?xi32>) -> tensor<?xi32> {
  %0 = mhlo.constant dense<0> : tensor<i32>
  %1 = mhlo.constant dense<6> : tensor<i32>
  %2 = "chlo.broadcast_minimum"(%arg0, %1) {broadcast_dimensions = array<i64>} : (tensor<?xi32>, tensor<i32>) -> tensor<?xi32>

// CHECK-DAG:   [[PARTIAL_INDEX:%.*]] = "tf.Fill"([[VECTOR_RANK]], [[ZERO_1]]) : (tensor<1xi32>, tensor<i32>) -> tensor<?xi32>
// CHECK-DAG:   [[EXPANDED_INDIX:%.*]] = "tf.ExpandDims"([[INDEX]], [[ZERO]]) : (tensor<i32>, tensor<i32>) -> tensor<1xi32>
// CHECK-DAG:   [[ITEM_INDEX:%.*]] = "tf.Concat"([[ZERO]], [[EXPANDED_INDIX]], [[PARTIAL_INDEX]]) : (tensor<i32>, tensor<1xi32>, tensor<?xi32>) -> tensor<?xi32>

  // CHECK-NEXT: return [[exec_result]] : tensor<*xi32>
  %0 = "tf.ReadVariableOp"(%arg1) {device = "/CPU:0"} : (tensor<*x!tf_type.resource>) -> tensor<*xi32>
  %dyn_arg0 = "tf.SetStaticDimensionBounds" (%arg0, %arg2) :(tensor<?x?xi32>, tensor<2xi64>) -> tensor<?x?xi32>
  %dyn_0 = "tf.SetStaticDimensionBounds" (%0, %arg3) :(tensor<*xi32>, tensor<?xi64>) -> tensor<?xi64>