}
    }

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

// 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>