// Create input_shapes.
    TF_ShapeAndTypeList* input_shapes =
        TF_NewShapeAndTypeList(input_shapes_vec.size());
    for (size_t i = 0; i < input_shapes_vec.size(); ++i) {
      const auto& input_shape = input_shapes_vec[i];
      if (input_shape.has_value()) {
        TF_ShapeAndTypeListSetShape(input_shapes, i, input_shape->data(),
                                    input_shape->size());

    Output input_shape_i = ops::Slice(
        host_scope.WithOpName("input_shape_", i), input_shape,
        constant_pool.Get1DHostConstant(i), constant_pool.Get1DHostConstant(1));

    slice_size.push_back(ops::Sub(host_scope.WithOpName("slice_size_", i),
                                  input_shape_i, begin_i));
    DCHECK_EQ(slice_size.back().type(), DT_INT64);
  }

  // Set input_shapes.
  for (int i = 0; i < num_inputs; ++i) {
    std::vector<DimensionHandle> dims;
    const TF_ShapeAndType& input_shape = input_shapes->items[i];
    if (input_shape.num_dims == InferenceContext::kUnknownRank) {
      c.SetInput(i, c.UnknownShape());
      continue;
    }
    dims.reserve(input_shape.num_dims);
    for (int j = 0; j < input_shape.num_dims; ++j) {

    const std::vector<std::optional<std::vector<int>>>& input_shapes,
    const std::vector<std::string>& output_arrays,
    const std::vector<std::string>& control_output_arrays,
    const GraphdefToMlirOptions& import_options, mlir::MLIRContext* context) {
  auto module_or = GraphdefToMlirImport(
      input, input_arrays, input_dtypes, input_shapes, output_arrays,
      control_output_arrays, import_options, context);

  for (const int i : reduced_axes) {
    if (i < 0) {
      reduced_dims_inputs.push_back(
          Gather(scope, input_shape, Add(scope, Size(scope, input_shape), i)));
    } else {
      reduced_dims_inputs.push_back(Gather(scope, input_shape, i));
    }
  }
  const Output reduced_dims = Stack(scope, reduced_dims_inputs);
  Tensor reduced_axes_tensor(

void HandleConv2DInput(TF::Conv2DOp conv2d, int64_t block_size) {
  auto input = conv2d.getInput();
  auto input_shape = mlir::cast<RankedTensorType>(input.getType()).getShape();
  SmallVector<int64_t, 4> transform_shape = {
      input_shape[0], input_shape[1] / block_size, input_shape[2] / block_size,
      input_shape[3] * block_size * block_size};
  auto transform_result_type =

// CHECK-LABEL: func private @batched_function
// CHECK: arg1
func.func private @batched_function(%arg0: tensor<1x3xf32>, %arg1: tensor<*x!tf_type.resource>) -> tensor<1x3xf32>
  attributes {tf._input_shapes = [#tf_type.shape<1x3>, #tf_type.shape<*>], tf.signature.is_stateful} {
  // CHECK: [[handle:%.*]] = "tf.VarHandleOp"()
  // CHECK: "tf.ReadVariableOp"([[handle]])

  TensorShape updates_shape({4});
  TensorShape input_shape({8});
  auto input = Placeholder(scope_, DT_FLOAT, Placeholder::Shape(input_shape));
  auto updates =
      Placeholder(scope_, DT_FLOAT, Placeholder::Shape(updates_shape));
  auto indices = Const(scope_, {{4}, {3}, {1}, {7}});
  auto y = ScatterNdNonAliasingAdd(scope_, input, indices, updates);
  RunTest({input, updates}, {input_shape, updates_shape}, {y}, {input_shape});
}

                                     Value &padding, int num_dims) {
  ShapedType input_shape = mlir::cast<ShapedType>(input.getType());
  SmallVector<int64_t> spatial_dims(num_dims - 2);
  absl::c_iota(spatial_dims, 1);
  bool has_dynamic_spatial_dim = absl::c_any_of(
      spatial_dims,
      [&input_shape](int64_t dim) { return input_shape.isDynamicDim(dim); });
  if (conv_padding.strref() == "SAME" && has_dynamic_spatial_dim) {

                                    StringAttr conv_padding,
                                    ArrayAttr explicit_paddings) {
  auto input_shape = mlir::cast<ShapedType>(input.getType());
  auto filter_shape = mlir::cast<ShapedType>(filter.getType());
  if (!input_shape.hasRank() || input_shape.getRank() != 4 ||
      !filter_shape.hasRank() || filter_shape.getRank() != 4) {
    emitError(loc, "input and filter are expected to be 4D tensors");