auto dynamic_key =
          rewriter.create<TF::_XlaCompileMlirPlaceholderProgramKeyOp>(
              func.getLoc(), /*program=*/result_type, llvm::ArrayRef<Value>{});

      auto recv_at_host = rewriter.create<TF::_XlaRecvAtHostOp>(
          func.getLoc(), op.getOperandTypes(), /*dynamic_key=*/dynamic_key,
          op.getSendKeyAttr(),
          /*device_ordinal=*/rewriter.getI64IntegerAttr(0),

        mlir::isa<quant::QuantizedType>(element_type)) {
      return;
    }
    op_builder->setInsertionPoint(const_op);
    Value scalar = op_builder->create<mhlo::ConstantOp>(
        const_op->getLoc(),
        DenseElementsAttr::get(
            RankedTensorType::get(/*shape=*/{}, element_type),
            splat_elements_attr.getSplatValue<Attribute>()));
    auto broadcast_dims = DenseIntElementsAttr::get(

    BlockArgument param = callee.getArgument(i);
    if (mlir::isa<TF::ResourceType>(getElementTypeOrSelf(param.getType()))) {
      removed_params.set(i);
      callee.getBody().addArgument(param.getType(), param.getLoc());
      param.replaceAllUsesWith(callee.getArguments().back());
      removed_params.push_back(false);
    }
  }
  // Remove old resource-type parameters.
  callee.getBody().front().eraseArguments(removed_params);

    auto fcast_op = rewriter.create<TF::CastOp>(
        preceding_sc_op->getLoc(), dq_arg_type.clone(rewriter.getF32Type()),
        preceding_sc_op.getResult());

    // Create a new AvgPool op with float type.
    TF::AvgPoolOp float_avg_pool_op = rewriter.create<TF::AvgPoolOp>(
        avg_pool_op->getLoc(),
        avg_pool_op.getType().clone(rewriter.getF32Type()),
        /*operands=*/fcast_op.getResult(),

    auto name_from_loc = mlir::GetNameFromLoc(op->getLoc());
    if (!name_from_loc.empty()) return name_from_loc;
    // If the location is none of the expected types, then simply use name
    // generated using the op type.
    return std::string(op->getName().getStringRef());
  }
  auto val = op_or_val.dyn_cast<mlir::Value>();
  auto name_from_loc = mlir::GetNameFromLoc(val.getLoc());
  if (!name_from_loc.empty()) return name_from_loc;

}

static Value CreateTranspose(OpBuilder& builder, Value source,
                             ArrayRef<int64_t> perm) {
  return builder.create<stablehlo::TransposeOp>(source.getLoc(), source, perm)
      ->getResult(0);
}

// Transform pad(transpose(x)) to transpose(pad(x))
struct TransposeCommuteWithPad : public OpRewritePattern<stablehlo::PadOp> {

  // Create new TFL While op that will be used to replace TF While op.
  auto new_op = OpBuilder(op).create<TFL::WhileOp>(
      op->getLoc(), op->getResultTypes(), op->getOperands(),
      while_op.getIsStateless());
  Location loc = while_op->getLoc();
  CreateRegionWithCall(while_op.cond_function(), new_op.getCond(), loc);
  CreateRegionWithCall(while_op.body_function(), new_op.getBody(), loc);

  auto collectOps = [this](Operation *op) {
    const auto &op_name = op->getName().getStringRef().str();
    if (absl::StartsWith(op_name, "tf.") || absl::StartsWith(op_name, "tfl.")) {
      loc_to_name_.emplace(op->getLoc(), op_name);
    }
  };

  for (auto &region : module->getRegions()) {
    region.walk(collectOps);
  }

  pass_name_ = extract_pass_name(pass->getName().str());
  error_collector_->Clear();
}

    for (int i = 0; i < op->getNumOperands(); i++) {
      Value input = op->getOperand(i);
      if (getElementTypeOrSelf(input).isBF16()) {
        Value f32_cast = rewriter.create<TF::CastOp>(
            op->getLoc(),
            CloneTypeWithNewElementType(input.getType(), rewriter.getF32Type()),
            input);
        op->setOperand(i, f32_cast);
      }
    }

    // Casts BF16 outputs of the operation.

    auto new_reshape_type = RankedTensorType::get(new_reshape_dims, el_ty);
    ReshapeOp new_reshape =
        rewriter.create<ReshapeOp>(new_reshape_shape.getLoc(), new_reshape_type,
                                   op.getInput(), new_reshape_shape);
    TF::ConstOp new_broadcast_shape =
        GetI64ConstantTensor(rewriter, reshape_shape, op.getLoc());