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

    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;

    zero = builder.getFloatAttr(elem_type, 0);
  } else {
    return false;
  }
  if (auto ranked_type = dyn_cast<RankedTensorType>(type)) {
    llvm::ArrayRef<int64_t> type_shape = ranked_type.getShape();
    for (int64_t i : type_shape) {
      if (i < 0) return false;
    }
    shape = builder.getI64TensorAttr(type_shape);
  } else {
    return false;
  }
  return true;
}

GetBroadcastShapesForBatchMatmul(ShapedType input_type,
                                 ShapedType weight_type) {
  ArrayRef<int64_t> input_shape = input_type.getShape();
  ArrayRef<int64_t> weight_shape = weight_type.getShape();

  const int64_t num_matmul_dim = 2;
  const int64_t num_input_batch_dim = input_type.getRank() - num_matmul_dim;

    auto input_type = mlir::dyn_cast<ShapedType>(op.getOperand(0).getType());

    llvm::ArrayRef<int64_t> weight_shape = weight_type.getShape();
    llvm::ArrayRef<int64_t> input_shape = input_type.getShape();

    // If weight_shape[2] != 1, it means weight shape was already restored.
    if (weight_shape[2] != 1) return failure();

    // Weight was reshaped into [H, W, 1, InxMul].