ASSERT_EQ(errors::OK, s.code()) << s.message();
      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
    enable_tensor_float_32_execution(false);
  }

  AbstractContextPtr ctx_;

 public:

    device_threads_.emplace_back(new DeviceThread(
        devices[device_index].c_str(), is_async, in_flight_nodes_limit));
  }
}

// Necessary for a unique_ptr to a forward-declared type.
ParallelDevice::~ParallelDevice() = default;

std::unique_ptr<ParallelTensor> ParallelDevice::CopyToParallelDevice(
    TFE_Context* context, TFE_TensorHandle* tensor, TF_Status* status) const {

    //   |      dz|          |
    //   |        |          |
    //   |     Const_3       |
    //   |                   |
    //   |        ^          |
    //   |       z|          |     // MatMul Forward Graph
    //   |        |          |
    //   |      MatMul       |
    //   |     /       \     |
    //   |    ^         ^    |
    //   |    |         |    |
    //   |---x|        y|----|