var_to_data_type_map_.swap(result.second);
  } else {
    reader_.reset(new TensorSliceReader(filename));
    if (!reader_->status().ok()) {
      tsl::Set_TF_Status_from_Status(status, reader_->status());
      return;
    }
    var_to_shape_map_.reset(
        new TensorSliceReader::VarToShapeMap(reader_->GetVariableToShapeMap()));
    var_to_data_type_map_.reset(new TensorSliceReader::VarToDataTypeMap(

      status_ =
          BuildImmediateExecutionContext(std::get<1>(GetParam()), &ctx_raw);
      ASSERT_EQ(errors::OK, status_.code()) << status_.message();
      immediate_execution_ctx_.reset(ctx_raw);
    }

    // Computing numerical gradients with TensorFloat-32 is numerically
    // unstable. Some forward pass tests also fail with TensorFloat-32 due to
    // low tolerances

  // TODO(b/172003047): Consider making AbstractOperation RefCounted.
  // parent_op->Unref();
}
Status TapeOperation::Reset(const char* op, const char* raw_device_name) {
  forward_op_.op_name = op;
  forward_op_.attrs.Reset(op);
  forward_op_.inputs.clear();
  forward_op_.outputs.clear();
  return parent_op_->Reset(op, raw_device_name);
}
const string& TapeOperation::Name() const { return parent_op_->Name(); }

class WorkerCacheInterface;

class ImmediateExecutionDistributedManager {
 public:
  virtual ~ImmediateExecutionDistributedManager() {}

  // Set up distributed execution environment on local and remote tasks.
  // When `reset_context` is true, initialize new cluster context state based
  // on cluster configurations provided in `server_def`; otherwise, update
  // existing context state with the provided `server_def`. Contexts created

    ASSERT_EQ(errors::OK, s.code()) << s.message();
    ctx.reset(ctx_raw);
  }

  AbstractTensorHandlePtr x;
  {
    AbstractTensorHandle* x_raw = nullptr;
    Status s = TestScalarTensorHandle<float, TF_FLOAT>(ctx.get(), 2.0f, &x_raw);
    ASSERT_EQ(errors::OK, s.code()) << s.message();
    x.reset(x_raw);
  }

  Status s = RunModel(TestScalarShape, ctx.get(),

    ASSERT_EQ(errors::OK, s.code()) << s.message();
    x.reset(x_raw);
  }

  AbstractTensorHandlePtr y;
  {
    AbstractTensorHandle* y_raw = nullptr;
    Status s =
        TestScalarTensorHandle<float, TF_FLOAT>(ctx_.get(), 7.0f, &y_raw);
    ASSERT_EQ(errors::OK, s.code()) << s.message();
    y.reset(y_raw);
  }

  float expected_dx[1] = {7.0f};

    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
    x1.reset(x1_raw);
  }

  AbstractTensorHandlePtr x2;
  {
    AbstractTensorHandle* x2_raw = nullptr;
    status_ = TestScalarTensorHandle<float, TF_FLOAT>(
        immediate_execution_ctx_.get(), 1.0f, &x2_raw);
    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
    x2.reset(x2_raw);
  }

                             ->Finalize(&output_list, &func));
      scoped_func.reset(func);
      for (auto output : output_list.outputs) {
        output->Unref();
      }
      TF_RETURN_IF_ERROR(ctx->RegisterFunction(func));
    }

    AbstractOperationPtr fn_op(ctx->CreateOperation());
    TF_RETURN_IF_ERROR(fn_op->Reset(fn_name, /*raw_device_name=*/nullptr));
    for (auto input : inputs) {

if [[ "${TFCI_MACOS_UPGRADE_PYENV_ENABLE}" == 1 ]]; then
  # The TFCI Mac VM image seems to have uncommitted local changes to the Pyenv
  # repository so we have to discard them and reset the working directory before
  # we can pull in the latest changes.
  cd /Users/kbuilder/.pyenv/ && git reset --hard HEAD && git pull && cd -
fi

# "TFCI_MACOS_PYENV_INSTALL_ENABLE" controls whether to use Pyenv to install

RUN groupadd -g 1001 buildslave && useradd -m -u 1001 -g buildslave buildslave
RUN mkdir -p /tf/venv
RUN chown -R buildslave:buildslave /tf
RUN dpkg -i /patchelf/patchelf_0.14.3-1_arm64.deb

# All lines past this point are reset when $CACHEBUSTER is set. We need this
# for Python specifically because we install some nightly packages which are
# likely to change daily.
ARG CACHEBUSTER=0
RUN echo $CACHEBUSTER

# Setup build and environment