std::vector<shape_inference::ShapeHandle> input_shapes;
  TF_CHECK_OK(c.input("input", &input_shapes));
  ASSERT_EQ("[3,4]", c.DebugString(input_shapes[0]));
  TF_CHECK_OK(reg->shape_inference_fn(&c));
  ASSERT_EQ("[3,4,8]", c.DebugString(c.output(0)));
}

TEST(BitcastOpTest, TestShapeInference_SmallerShape) {
  const OpRegistrationData* reg;

  if (!status->status.ok()) return;

  if (op_reg_data->shape_inference_fn == nullptr) {
    status->status =
        InvalidArgument("No shape inference function exists for op '",
                        node_def.op(), "', did you forget to define it?");
    return;
  }

  status->status = c.Run(op_reg_data->shape_inference_fn);
  if (!status->status.ok()) return;

  // Set output_shapes.

  if (!op_reg_data) {
    DVLOG(1) << "Skipping inference for unregistered op " << node.type_string();
    return default_type();
  }
  if (op_reg_data->shape_inference_fn == nullptr) {
    DVLOG(1) << "Skipping inference for op without shape function "
             << node.type_string();
    return default_type();
  }

}

void TF_OpDefinitionBuilderSetShapeInferenceFunction(
    TF_OpDefinitionBuilder* builder,
    void (*shape_inference_func)(TF_ShapeInferenceContext* ctx,
                                 TF_Status* status)) {
  auto* cc_builder = reinterpret_cast<OpDefBuilder*>(builder);
  cc_builder->SetShapeFn(
      [shape_inference_func](InferenceContext* ctx) -> tensorflow::Status {
        TF_Status* c_status = TF_NewStatus();

// Sets the shape inference function for the op.
TF_CAPI_EXPORT extern void TF_OpDefinitionBuilderSetShapeInferenceFunction(
    TF_OpDefinitionBuilder* builder,
    void (*shape_inference_func)(TF_ShapeInferenceContext* ctx,
                                 TF_Status* status));

//----------------------------------------------------
// Functions for TF_ShapeInferenceContext.
//