CHECK_EQ(TF_OK, TF_GetCode(status_)) << TF_Message(status_);
  }

  ~ShapeInferenceTest() override {
    TFE_DeleteContextOptions(tfe_context_options_);
    TFE_DeleteContext(tfe_context_);
    TF_DeleteStatus(status_);
  }

  // Checks the expected result of shape inference for the given `op`.
  void CheckOutputShapes(
      TFE_Op* op,