}
    }

    // Sometimes we have VariableShape nodes in while loop (after Enter nodes).
    // They won't be constant-folded because TensorFlow constant folding does
    // not handle Enter nodes (and thus does not handle any nodes after Enter
    // nodes). We try to replace such VariableShape nodes with Const nodes here.
    if (n->type_string() == "VariableShape") {

    tf_device.return
  }) : () -> ()
  func.return
}

// Tests that tf.VariableShape operation is decomposed.

// CHECK-LABEL: @decompose_variable_shape_i32
func.func @decompose_variable_shape_i32(%input: tensor<!tf_type.resource<tensor<?x?x?xf32>>>) -> tensor<3xi32> {
  %0 = "tf_device.cluster"() ({
    %1 = "tf.VariableShape"(%input) : (tensor<!tf_type.resource<tensor<?x?x?xf32>>>) -> tensor<3xi32>

      (TF_ExpandDimsOp $indices,
       (TF_ConstOp (GetScalarOfType<-1> $indices))),
      $updates
    ),
    (CreateConstBoolAttrFalse)
  )>;

// Pattern to decompose tf.VariableShape into tf.ReadVariable and tf.Shape.
def DecomposeVariableShape : Pat<
  (TF_VariableShapeOp:$src_op $resource),
  (TF_ShapeOp (CreateTFReadVariableOpFromResourceHandle $src_op, $resource)),

    ```mlir
    tf_device.replicate([%0, %1] as %ri: tensor<*x!tf_type.resource>) {n = 2 : i32} {
      %2 = "tf.ReadVariableOp"(%ri) : tensor<*x!tf_type.resource> -> tensor<*xi32>
      %3 = "tf.VariableShape"(%0) : (tensor<*x!tf_type.resource>) -> tensor<?xi32>
      tf_device.return
    }
    ```
  }];

  let dependentDialects = [
    "mhlo::MhloDialect",
    "tf_device::TensorFlowDeviceDialect"

//===----------------------------------------------------------------------===//
// ShapeOp
//===----------------------------------------------------------------------===//

namespace {
// Validates Shape/ShapeN/VariableShape operand and associated result types.
LogicalResult VerifyShapeOperandAndResult(Operation *op, Type operand_type,
                                          Type result_type,

  let results = (outs
    TF_Tensor:$ref
  );

  TF_DerivedResultTypeAttr dtype = TF_DerivedResultTypeAttr<0>;

  let hasCanonicalizer = 1;
}

def TF_VariableShapeOp : TF_Op<"VariableShape", []> {
  let summary = "Returns the shape of the variable pointed to by `resource`.";

  let description = [{
This operation returns a 1-D integer tensor representing the shape of `input`.