// TODO(ycao): Support computation with compile-time constant, which requires
// non-trivial input mapping as implemented now.
void GetInputMappingForMlir(int num_inputs, std::vector<int>* input_mapping) {
  input_mapping->resize(num_inputs, 0);
  std::iota(input_mapping->begin(), input_mapping->end(), 0);
}

static void RegisterDialects(mlir::DialectRegistry& registry) {

    return errors::InvalidArgument("Equation must contain a single ->");
  }

  const absl::string_view input_subs = equation_split[0];
  const absl::string_view output_subs = equation_split[1];
  if (op.num_inputs() == 1) {
    // For the unary einsum z = einsum("{eq_x}->{eq_z}", x), the gradient wrt
    // the input (VJP) is given by the reversed equation:
    //   grad_wrt_x = einsum("{eq_z}->{eq_x}", grad_wrt_z)

TF_CAPI_EXPORT extern void TFE_OpAddInputList(TFE_Op* op,
                                              TFE_TensorHandle** inputs,
                                              int num_inputs,
                                              TF_Status* status);

// Fetches the current number of inputs attached to `op`.
//
// Does not use the operation's definition to determine how many inputs should

  }
}

bool TF_IsRefInput(TF_OpKernelContext* ctx, int i, TF_Status* status) {
  auto* cc_ctx = reinterpret_cast<::tensorflow::OpKernelContext*>(ctx);
  if (i < 0 || i >= cc_ctx->num_inputs()) {
    TF_SetStatus(status, TF_OUT_OF_RANGE, "input index out of range");
    return false;
  }
  TF_SetStatus(status, TF_OK, "");
  return cc_ctx->input_is_ref(i);
}

#ifndef IS_MOBILE_PLATFORM

      send_recv_nodes.push_back(n);
    }
  }
  EXPECT_EQ(num_send_from_host, 1);
  EXPECT_EQ(num_recv_at_host, 1);
  for (Node *n : send_recv_nodes) {
    Node *input_node;
    TF_CHECK_OK(n->input_node(n->num_inputs() - 1, &input_node));
    EXPECT_EQ(input_node, key_placeholder);

    bool has_control_edge_to_sequencer = false;
    for (const Edge *e : n->out_edges()) {
      if (e->IsControlEdge() && e->dst() == sequencer) {

    : OpKernel(ctx), platform_info_(XlaPlatformInfoFromDevice(ctx->device())) {}

void XlaRunOp::Compute(OpKernelContext* ctx) {
  VLOG(3) << "XlaRunOp " << def().name();
  Tensor key_tensor = ctx->input(ctx->num_inputs() - 1);

  bool use_pjrt =
      GetXlaOpsCommonFlags()
          ->tf_xla_use_device_api.IsEnabledInXlaCompileAndRunForDevice(
              platform_info_.device_type());

  if (use_pjrt) {

  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  TFE_OpAddAttrs(ret, TFE_OpGetAttrs(other));
  int num_inputs = TFE_OpGetFlatInputCount(other, status);
  CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
  for (int input_index = 0; input_index < num_inputs; ++input_index) {
    TFE_TensorHandle* input = TFE_OpGetFlatInput(other, input_index, status);

// For inputs that take a list of tensors.
// inputs must point to TF_Output[num_inputs].
TF_CAPI_EXPORT extern void TF_AddInputList(TF_OperationDescription* desc,
                                           const TF_Output* inputs,
                                           int num_inputs);

// Call once per control input to `desc`.

    mlir::FunctionType callee_type = callee.getFunctionType();
    SmallVector<Type, 4> expected_callee_input_types;
    auto num_inputs = op.getInputs().size() / 2;
    for (unsigned i = 0; i < num_inputs; ++i) {
      auto input_type =
          mlir::dyn_cast<RankedTensorType>(op.getOperand(i).getType());
      auto scalar = RankedTensorType::get({}, input_type.getElementType());

                                        ? mhlo::Precision::DEFAULT
                                        : mlir::mhlo::Precision::HIGHEST;
  llvm::SmallVector<mlir::Attribute, 2> attr_vec;
  const int num_inputs = 2;
  for (int i = 0; i < num_inputs; i++) {
    attr_vec.push_back(
        mlir::mhlo::PrecisionAttr::get(builder->getContext(), precision));
  }
  return builder->getArrayAttr(attr_vec);
}