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