// CHECK-LABEL: func private @batch_0
// CHECK: f = @compute_0
func.func private @batch_0(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {

// the batch functions with same shared_name.

func.func private @batch_0(%arg0: tensor<?x?xi32>) -> tensor<*xi32> {

  "tf.TPUCompileSucceededAssert"(%compilation_status) {device = "/CPU:0"} : (tensor<!tf_type.string>) -> ()
  %0 = "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/CPU:0", shared_name = "y"} : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>

  Graph g(OpRegistry::Global());
  TF_CHECK_OK(s.ToGraph(&g));

  TF_CHECK_OK(PerformStaticShapeInferenceBeforeEncapsulation(&g));

  // Check that "add" node now has _xla_inferred_shapes attr.
  auto node_index = g.BuildNodeNameIndex();
  Node *add_node = node_index["add"];
  std::vector<PartialTensorShape> output_shapes;
  TF_CHECK_OK(GetNodeAttr(add_node->attrs(), kXlaInferredShapesAttrName,

      %cst = "tf.Const"() {value = dense<[1, -1, 4]> : tensor<3xi32>} : () -> tensor<3xi32>
      %0 = "tf.Slice"(%arg0, %arg1, %cst) {_XlaHasReferenceVars = false, _xla_inferred_shapes = [#tf_type.shape<1x?x4>], device = "/job:localhost/replica:0/task:0/device:TPU:0"} : (tensor<3x3x10xbf16>, tensor<3xi32>, tensor<3xi32>) -> tensor<1x?x4xbf16>
      return %0 : tensor<1x?x4xbf16>
    }
  })";

        }
        g->AddControlEdge(iter->second, n);
      }
    }
  }
  return absl::OkStatus();
}
}  // namespace

const char kXlaInferredShapesAttrName[] = "_xla_inferred_shapes";

const char kXlaConnectedToXlaComputationAttrName[] =
    "_xla_connected_to_xla_computation";
const char kXlaConnectedFromXlaComputationAttrName[] =
    "_xla_connected_from_xla_computation";

  graphdef->Swap(&graphdef_out);

  *library = lib_def->ToProto();
  // Remove "_xla_inferred_shapes" attr. They are added by
  // `PerformStaticShapeInferenceBeforeEncapsulation`.
  for (FunctionDef& fdef : *library->mutable_function()) {
    for (NodeDef& node_def : *fdef.mutable_node_def()) {
      node_def.mutable_attr()->erase("_xla_inferred_shapes");
    }
  }

  return s;
}

    // CHECK-NOT: tf.ResourceApplyAdam

  %3 = "tf.AddV2"(%arg0, %1) {_xla_inferred_shapes = [#tf_type.shape<>], device = ""} : (tensor<i32>, tensor<i32>) -> tensor<i32>
  %4 = "tf.Pack"(%3) {_xla_inferred_shapes = [#tf_type.shape<1>], axis = 0 : i64, device = ""} : (tensor<i32>) -> tensor<1xi32>
  %5 = "tf.Pack"(%arg0) {_xla_inferred_shapes = [#tf_type.shape<1>], axis = 0 : i64, device = ""} : (tensor<i32>) -> tensor<1xi32>

  %outputs_8  =  "tf.Const"() {device = "/job:localhost/replica:0/task:0/device:CPU:0", value = dense<1> : tensor<i32>} : () -> tensor<i32>
  // CHECK: [[elems:%.*]] = "tf.VarHandleOp"
  %outputs_10  =  "tf.VarHandleOp"() {_xla_inferred_shapes = [#tf_type.shape<>], allowed_devices = [], container = "", device = "/job:localhost/replica:0/task:0/device:CPU:0", shared_name = "w"} : () -> tensor<!tf_type.resource<tensor<3x1xf32>>>