string func_name;
    string node_name;
    OP_REQUIRES_OK(ctx, ctx->GetAttr("log_dir_path", &log_dir_path));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("enabled", &enabled_));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("file_name", &file_name));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("func_name", &func_name));
    OP_REQUIRES_OK(ctx, ctx->GetAttr("node_name", &node_name));
    OP_REQUIRES_OK(ctx, ctx->env()->RecursivelyCreateDir(log_dir_path));

      : OpKernel(context) {
    OP_REQUIRES_OK(context, context->GetAttr("id", &id_));

    int calibration_method_value;
    int num_bins;
    float min_percentile;
    float max_percentile;
    OP_REQUIRES_OK(context, context->GetAttr("calibration_method",
                                             &calibration_method_value));
    OP_REQUIRES_OK(context, context->GetAttr("num_bins", &num_bins));

    std::string output_file_path;
    OP_REQUIRES_OK(context,
                   context->GetAttr("output_file_path", &output_file_path));
    OP_REQUIRES_OK(context, context->env()->NewWritableFile(output_file_path,
                                                            &output_file_));

    OP_REQUIRES_OK(context, context->GetAttr("ids", &ids_));
    OP_REQUIRES_OK(context, context->GetAttr("calibration_methods",

    // Lock variables for the whole duration of compile + execute.
    OP_REQUIRES_OK(ctx, GetAndLockVariablesAndBuildXlaCompilerArguments(
                            *ctx, inputs, *constant_indices_or,
                            variable_indices, &variables, &args));

    PjRtDeviceCompiler* pjrt_device_compiler;
    xla::PjRtLoadedExecutable* pjrt_executable;
    OP_REQUIRES_OK(ctx, Compile(args, ctx, &pjrt_device_compiler, &profiler,

  OP_REQUIRES_OK(context, s);
  allocFunc(ctx, tmp_var_tf.get(), dtype, dims, num_dims, tf_status);
  s = tensorflow::StatusFromTF_Status(tf_status);
  if (!s.ok()) tmp_var->Unref();
  OP_REQUIRES_OK(context, s);

  OP_REQUIRES_OK(context, TF_TensorToTensor(tmp_var_tf.get(), &tmp_var->val));
  OP_REQUIRES_OK(context,

          GatherVariableInfo(ctx, *closure.compilation_result(),
                             closure.num_constant_args());
      OP_REQUIRES_OK(ctx, updated_variables.status());
      OP_REQUIRES_OK(ctx, LockVariables(absl::MakeSpan(*updated_variables)));
      OP_REQUIRES_OK(
          ctx, RunPjRtExecutable(closure.num_constant_args(), inputs,
                                 variable_snapshots, *updated_variables,

             << type_string() << " on an XLA device. This should never happen.";
}

XlaAssignVariableOp::XlaAssignVariableOp(OpKernelConstruction* c)
    : OpKernel(c) {
  OP_REQUIRES_OK(c, c->GetAttr("dtype", &dtype_));
}

void XlaAssignVariableOp::Compute(OpKernelContext* context) {
  OP_REQUIRES(context, dtype_ == context->input(1).dtype(),
              errors::InvalidArgument(

  void Compile(XlaOpKernelContext* context) override;

 private:
  string name_;
};

XlaAotOnlyVarHandleOp::XlaAotOnlyVarHandleOp(OpKernelConstruction* c)
    : XlaOpKernel(c) {
  OP_REQUIRES_OK(c, c->GetAttr("shared_name", &name_));
}

void XlaAotOnlyVarHandleOp::Compile(XlaOpKernelContext* context) {
  // Look for a resource of the same name. TF also keys that on the container