}
    }

    // 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") {

                              .TypeConstraint("T", TYPES),                     \
                          ShapeNOp<int64_t>);                                  \
  REGISTER_KERNEL_BUILDER(Name("VariableShape")                                \
                              .Device(DEVICE)                                  \
                              .TypeConstraint<int32>("out_type")               \

    tf_device.return %1, %2, %3 : tensor<*xf32>, tensor<?xi32>, tensor<*xi32>
  }
  func.return
}

// CHECK: %[[VAR_SHAPE:[0-9]*]] = "tf.VariableShape"(%[[ARG_0]])
// CHECK: tf_device.replicate([%[[ARG_0]], %[[ARG_1]]] as %[[RI:[a-z0-9]*]]: tensor<*x!tf_type.resource>)
// CHECK:   %[[READ_VAR:[0-9]*]] = "tf.ReadVariableOp"(%[[RI]])

    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)),

  %0 = "tf.VariableShape"(%arg0) {output = "tfdtype$DT_INT32"} : (tensor<*x!tf_type.resource<tensor<1x32x32x16xf32>>>) -> tensor<4xi32>
  %1 = "tf.VariableShape"(%arg1) {output = "tfdtype$DT_INT32"} : (tensor<*x!tf_type.resource>) -> tensor<?xi32>
  func.return %0, %1 : tensor<4xi32>, tensor<?xi32>
}

// Tests ops that variable shapes are correctly evaluated on static types.
// CHECK-LABEL: func @testVariableShape
func.func @testVariableShape(%arg0: tensor<!tf_type.resource<tensor<2x4xf32>>>) -> tensor<2xi32> {
  %0 = "tf.VariableShape"(%arg0) : (tensor<!tf_type.resource<tensor<2x4xf32>>>) -> tensor<2xi32>
  // CHECK:         [[cst:%.*]] = "tf.Const{{.*}} dense<{{\[}}2, 4]> : tensor<2xi32>
  // CHECK-NEXT:    return [[cst]] : tensor<2xi32>

      "UniqueV2",
      "UpperBound",
      "UnsortedSegmentMax",
      "UnsortedSegmentMin",
      "UnsortedSegmentProd",
      "UnsortedSegmentSum",
      "VarIsInitializedOp",
      "VariableShape",
      "Where",
      "While",
      "XlaAllReduce",
      "XlaBroadcastHelper",
      "XlaCallModule",
      "XlaConcatND",
      "XlaConv",
      "XlaConvV2",
      "XlaCustomCall",

    ```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"

```

gets converted to

```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
}
```
### `-tf-replicate-tensor-list-init-ops`

_Replicate TensorList init ops for correct shape assignments in shape inference_