// CHECK: name: "StatefulWhile"
// CHECK-NOT: name:
// CHECK: op: "While"
// CHECK-NOT: is_stateless
// CHECK-NOT: shape_invariant
// CHECK:  attr {
// CHECK:    key: "output_shapes"
// CHECK:    value {
// CHECK:      list {
// CHECK:        shape {
// CHECK:          dim {
// CHECK:            size: 5
// CHECK:          }
// CHECK:        }
// CHECK:      }
// CHECK:    }

  %4 = "tf.TensorSliceDataset"(%2) {device = "", output_shapes = [#tf_type.shape<>], metadata = ""} : (tensor<3xi32>) -> tensor<*x!tf_type.variant>
  // CHECK: "tf.MapAndBatchDataset"(%[[TSLICE]],
  // CHECK-SAME: f = @"__inference_Dataset_map_<lambda>_80",
  %5 = "tf.ParallelMapDataset"(%4, %3) {device = "",
           f = @"__inference_Dataset_map_<lambda>_80",
           output_shapes = [#tf_type.shape<>], output_types = [i32],

template <typename Op>
struct DropAttributes : public OpRewritePattern<Op> {
  using OpRewritePattern<Op>::OpRewritePattern;

  // Drop the "output_shapes" attribute.
  LogicalResult matchAndRewrite(Op op,
                                PatternRewriter &rewriter) const override {
    bool found = !!op->removeAttr("output_shapes");
    return success(found);
  }
};

// Helper function to create TF op while copying all underscore attributes from

    %drop_remainder: i1 {tfr.name="drop_remainder"},
    %f: !tfr.attr {tfr.name="func"},
    %output_types: !tfr.attr {tfr.name="output_types"},
    %output_shapes: !tfr.attr {tfr.name="output_shapes"},
    %preserve_cardinality: i1 {tfr.name="preserve_cardinality", tfr.default=false}) -> !tfr.tensor {
  %batch = "tfr.constant_tensor"(%batch_size) : (i64) -> tensor<i64>

    ) {
    // CHECK: tf.ReduceDataset
    %1 = "tf.ReduceDataset"(%arg0, %arg1) {
      Targuments = [],
      Tstate = [i64], device = "",
      f = @__reduce_func_0, f._tf_data_function = true,
      output_shapes = [#tf_type.shape<>],
      output_types = [i64], use_inter_op_parallelism = true } : (tensor<!tf_type.variant>, tensor<i64>) -> (tensor<i64>)
    func.return
}

  op: "TensorSliceDataset"
  input: "tensors/normalize_tensors/component_0"
  attr {
    key: "Toutput_types"
    value {
      list {
        type: DT_INT32
      }
    }
  }
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
        }
      }
    }
  }
}
node {
  name: "MapDataset"
  op: "MapDataset"
  input: "TensorSliceDataset"
  attr {
    key: "Targuments"

  std::array<int64_t, 4> output_shape;
  // NHWC <- NCHW
  output_shape[0] = composite_result_shape[0];
  output_shape[1] = composite_result_shape[2];
  output_shape[2] = composite_result_shape[3];
  output_shape[3] = composite_result_shape[1];

  auto input_type = mlir::cast<ShapedType>(old_op->getOperand(0).getType());

  return RankedTensorType::get(output_shape, input_type.getElementType());
}

    CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
    CHECK_EQ(output_shapes->num_items, 1);

    int num_dims = output_shapes->items[0].num_dims;
    int64_t* dims = output_shapes->items[0].dims;

    if (!expected_shape.has_value()) {
      EXPECT_EQ(num_dims, -1);
      EXPECT_EQ(dims, nullptr);
      return;
    }

    EXPECT_EQ(num_dims, expected_shape->size());

    }
  }
  attr {
    key: "index"
    value {
      i: 4
    }
  }
}
node {
  name: "IteratorGetNext"
  op: "IteratorGetNext"
  input: "args_0"
  attr {
    key: "output_shapes"
    value {
      list {
        shape {
          dim {
            size: 1
          }
          dim {
            size: 8
          }
        }
        shape {
          dim {

      %outputs_4, %control_5 = tf_executor.island wraps "tf.Identity"(%outputs_2) {device = ""} : (tensor<*xi32>) -> tensor<*xi32>