}

        shape_inference::ShapeHandle handle;
        TF_RETURN_IF_ERROR(
            context->MakeShapeFromPartialTensorShape(arg_shape.shape, &handle));
        TF_RETURN_IF_ERROR(shape_refiner->SetShape(n, 0, handle));
      }
    }

    // Sometimes we have VariableShape nodes in while loop (after Enter nodes).
    // They won't be constant-folded because TensorFlow constant folding does

  // CHECK-NEXT: %[[BROADCAST:.*]] = "tf.BroadcastTo"(%[[CAST_ZERO]], %[[CONST10]]) : (tensor<f32>, tensor<1xi32>) -> tensor<10xf32>
  // CHECK-NEXT: %[[SIZE_SHAPE:.*]] = "tf.Const"() <{value = dense<1> : tensor<1xi32>}>
  // CHECK-NEXT: %[[SIZE:.*]] = "tf.Reshape"(%[[NUM]], %[[SIZE_SHAPE]])
  %tl = "tf.TensorListReserve"(%elem_shape, %num) : (tensor<0xi32>, tensor<i32>) -> tensor<!tf_type.variant<tensor<f32>>>

    if (auto ranked_tensor_type =
            mlir::dyn_cast<RankedTensorType>(operand_type)) {
      tensorflow::TensorShapeProto shape_proto;
      ConvertToTensorShapeProto(ranked_tensor_type.getShape(), &shape_proto);
      *arg->mutable_shape() = std::move(shape_proto);
    } else {
      arg->mutable_shape()->set_unknown_rank(true);
    }

    if (failed(SetOpSharding(op, input_shardings.getValue()[index],

    // CHECK: %[[RESHAPE0:.*]] = mhlo.reshape %[[BROADCAST0]] : (tensor<4x4x1xf32>) -> tensor<4x4xf32>
    // CHECK: %[[UPDATED_ARG0:.*]] = "mhlo.broadcast_in_dim"(%[[RESHAPE0]]) <{broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>}> : (tensor<4x4xf32>) -> tensor<4x4x4xf32>

    // CHECK: %[[RESHAPE1:.*]] = mhlo.reshape %arg1 : (tensor<4x1x4xf32>) -> tensor<4x4xf32>

namespace tensorflow {
namespace tfrt_compiler {
namespace {

constexpr int64_t kDefaultCheapCost = 1;

int64_t GetRankedTensorSize(mlir::TensorType type) {
  auto shape = type.getShape();

  int64_t size = 1;
  for (int64_t dim : shape) {
    // For unknown dimensions, use 1 as the size because it is usually the batch
    // dimension.
    //

                                                    ArrayRef<int64_t> perm,
                                                    bool isInvert) {
  SmallVector<int64_t, 4> permutedShape = applyPermutation(
      type.getShape(), isInvert ? invertPermutationVector(perm) : perm);
  return RankedTensorType::get(permutedShape, type.getElementType());
}

static RankedTensorType GetInvertPermutedTensorType(RankedTensorType type,

    const auto value_attr =
        mlir::cast<DenseFPElementsAttr>(const_op.getValue());
    const ArrayRef<int64_t> original_shape =
        value_attr.getShapedType().getShape();

    const SmallVector<float> original_values =
        llvm::to_vector(value_attr.getValues<float>());

    // Fold the constant value by transposing the values according to the

// CHECK-LABEL: dynamicReshape
// CHECK-NEXT:  %[[reshape:.*]] = "tfl.reshape"(%arg0, %arg1) : (tensor<*xf32>, tensor<2xi32>) -> tensor<?x?xf32>
// CHECK-NEXT:  return %[[reshape]] : tensor<?x?xf32>
}

func.func @dynamicReshapeI64(%arg0: tensor<*xf32>, %arg1: tensor<2xi64>) -> tensor<?x?xf32> {
  %0 = "tf.Reshape"(%arg0, %arg1) : (tensor<*xf32>, tensor<2xi64>) -> tensor<?x?xf32>
  func.return %0 : tensor<?x?xf32>

    `tf.nn.space_to_depth`.

    ```python
    images = tf.reshape(images, [batch, h // block_size, block_size,
                               w // block_size, block_size, c])
    images = tf.transpose(images, [0, 1, 3, 2, 4, 5])
    images = tf.reshape(images, [batch, h // block_size, w // block_size,
                               c * (block_size ** 2)])
    ```

    return true;
  }
};
std::unique_ptr<TargetHardwareOperation> CreateArithmeticOp() {
  return std::make_unique<GpuArithmeticOp>();
}

// Currently used for these ops:
// tfl.concatenation / tfl.reshape
class GpuConcatOp : public TargetHardwareOperation {
  double GetOpCost(mlir::Operation* op) const override {
    int64_t count;
    if (ArithmeticCountUtilHelper::GetInputTensorTotalSize(op, &count))