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

    // TODO(hinsu): Handle constants that are not scalar booleans.
    auto cond_type = mlir::dyn_cast<RankedTensorType>(cond.getType());
    if (!cond_type || !cond_type.getShape().equals({}) ||
        !cond_type.getElementType().isInteger(/*width=*/1))
      return failure();

    // Identify the branch to inline.
    bool cond_value = (*cond.value_begin<APInt>()).getSExtValue();

    // Only fuses multiplier if all dimensions other than the out channel
    // dimension are equal to 1.
    if (!TFL::IsDimensionsDegenerateExceptLastOne(
            mul_value.getShapedType().getShape())) {
      return rewriter.notifyMatchFailure(mul_op, [&](::mlir::Diagnostic &diag) {
        diag << "entities 'mul_value' failed to satisfy constraint: "
                "unsupported dimensions";
      });
    }

  Type dtype = getElementTypeOrSelf(type);
  if (RankedTensorType ty = type.dyn_cast<RankedTensorType>()) {
    llvm::SmallVector<int64_t, 4> shape = {dim};
    shape.append(ty.getShape().begin(), ty.getShape().end());
    return tensorflow::GetTypeFromTFTensorShape(shape, dtype);
  }
  return type;
}

Type GetTensorTypeForTensorList(Type element_type, TF::VariantType handle_dtype,

    bool same_element_type = output_tensor_type.getElementType() ==
                             input_vector_type.getElementType();
    bool same_shape =
        output_tensor_type.getShape() == input_vector_type.getShape();
    if (!same_element_type || !same_shape) {
      op.emitError("input and output should have same shape and element type.");
    }
    return success(same_element_type && same_shape);
  }

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

    // Connect `expand_op` with the input of `stb_op`.
    expand_op.setOperand(0, stb_op.getInput());
    // Calculate the shape for expand.
    auto input_shape =
        mlir::cast<ShapedType>(stb_op.getInput().getType()).getShape();
    SmallVector<int64_t, 4> expand_shape(input_shape.begin(),
                                         input_shape.end());
    expand_shape.insert(expand_shape.begin() + expand_axis, 1);

    const auto result_type = mlir::cast<TensorType>(op.getResult(0).getType());
    const SmallVector<int64_t> new_result_shape =
        Permute<int64_t>(result_type.getShape(), kNchwToNhwcPermutation);

    const TensorType new_result_type =
        result_type.cloneWith(new_result_shape, result_type.getElementType());

                                  DenseIntElementsAttr bcast_dims) {
  auto dimensions = llvm::to_vector(bcast_dims.getValues<int64_t>());
  const auto result_shape = result_type.getShape();
  // Index for the broadcasted matrix.
  llvm::SmallVector<int64_t, 16> current_index(result_type.getRank(), 0);
  // Computes the new operand shape using the original shape and the broadcast

  if (!variant_type || variant_type.getSubtypes().size() != 1) return failure();
  TensorType tensor_type = variant_type.getSubtypes().front();
  if (!tensor_type.hasStaticShape()) return failure();
  for (auto d : tensor_type.getShape()) shape->push_back(d);
  return success();
}

LogicalResult HandleEmptyTensorListOp(
    TF::EmptyTensorListOp list,
    llvm::SmallDenseMap<Value, SizeInfo>* buffer_to_size) {
  Value buffer;