[Good First Issue][TF FE]: Support Case operation for TensorFlow models

[rkazants]

Context

OpenVINO component responsible for support of TensorFlow models is called as TensorFlow Frontend (TF FE). TF FE converts a model represented in TensorFlow opset to a model in OpenVINO opset.

In order to infer TensorFlow models with Case operation by OpenVINO, TF FE needs to be extended with this operation support.

What needs to be done?

For Case operation support, you need to add the corresponding loader into TF FE op directory and to register it into the dictionary of Loaders. One loader is responsible for conversion (or decomposition) of one type of TensorFlow operation.

Here is an example of loader implementation for TensorFlow Einsum operation:

OutputVector translate_einsum_op(const NodeContext& node) { 
     auto op_type = node.get_op_type(); 
     TENSORFLOW_OP_VALIDATION(node, op_type == "Einsum", "Internal error: incorrect usage of translate_einsum_op."); 
     auto equation = node.get_attribute<std::string>("equation"); 
  
     OutputVector inputs; 
     for (size_t input_ind = 0; input_ind < node.get_input_size(); ++input_ind) { 
         inputs.push_back(node.get_input(input_ind)); 
     } 
  
     auto einsum = make_shared<Einsum>(inputs, equation); 
     set_node_name(node.get_name(), einsum); 
     return {einsum}; 
 } 

In this example, translate_einsum_op converts TF Einsum into OV Einsum. NodeContext object passed into the loader packs all info about inputs and attributes of Einsum operation. The loader retrieves an attribute of the equation by using the NodeContext::get_attribute() method, prepares input vector, creates Einsum operation from OV opset and returns a vector of outputs.

Responsibility of a loader is to parse operation attributes, prepare inputs and express TF operation via OV operations sub-graph. Example for Einsum demonstrates the resulted sub-graph with one operation. In PR #19007 you can see operation decomposition into multiple node sub-graph.

Hint

Case operation is a type of body graph operations that needs special attention to extract and handle body graphs for each case. You can take a look how to approach with body graphs in loaders for While and If operations.

Example Pull Requests

• [TF FE][TF Hub] Support DivNoNan and EnsureShape operations #19007

Resources

• Contribution guide - start here!
• What is OpenVINO?
• OpenVINO TensorFlow Frontend
• Blog post on contributing to OpenVINO
• User documentation

Contact points

@openvinotoolkit/openvino-tf-frontend-maintainers

Ticket

TBD

[rghvsh]

Hey! @rkazants I'd like to work on this project

[rghvsh]

.take

[github-actions]

Thank you for looking into this issue! Please let us know if you have any questions or require any help.

[rghvsh]

Hi! @rkazants how to use a switch , should I use this
src/frontends/tensorflow_common/include/helper_ops/switch.hpp
src/frontends/tensorflow/src/op/switch.cpp

or just simply use the keyword switch

Thank You
Raghav

[rghvsh]

Hi! @rkazants
The Case operation has an n-way switch statement as n is variable how will we implement it(how will we have a variable number of cases).

Thank You
Raghav

[popovaan]

Hi @rghvsh!

The Case operation has an n-way switch statement as n is variable how will we implement it(how will we have a variable number of cases).

I would recommend you to start by creating a single layer TF model with Case operation tf.raw_ops.Case(...) (TF Case) and check how it looks like (using Netron for example).
Here is example of single-layer graph with Case:

We can see here that branches are defined as a list of TF functions in the attribute branches. This list has a definite size, so you don't have a variable N at this point. So in the translator of Case operation you need to create a body graph for each of these branches. Examples of creating body graphs you can find here.

how to use a switch , should I use this
src/frontends/tensorflow_common/include/helper_ops/switch.hpp
src/frontends/tensorflow/src/op/switch.cpp
or just simply use the keyword switch

Switch and Merge are helper operations, which are merged to OpenVino If operation. I would recommend to use If directly (Please check example of creation of If operation and If specification)

[rghvsh]

Hello @popovaan! I'll try this and get back to you.
Thanks for helping

Raghav

[p-wysocki]

Just yesterday our CONTRIBUTING.md has been updated with a technical guide - I highly recommend checking it out. :)

[rghvsh]

Hello! @popovaan sorry for the delay was not able to work in this had college practical exams. Have started working on this

Thanks
Raghav

[rkazants]

Hi @rghvsh,

Do you have any update? Please take a look at If operation support implementation here: https://github.com/openvinotoolkit/openvino/blob/master/src/frontends/tensorflow/src/op/if.cpp.

Case is just a general case of If with several condition values/bodies.

Best regards,
Roman

[rghvsh]

Hey @rkazants!
I have started working on this and made some progress will open a PR soon. Sorry for the delay my university theory exams session has started.

[rkazants]

Hi @rghvsh, is there update on this task? Please share your current code in PR and ask questions if you have.

Best regards,
Roman

[p-wysocki]

I am happy to announce that we have created a channel dedicated to Good First Issues support on our Intel DevHub Discord server! Join it to receive support, engage in discussions, ask questions and talk to OpenVINO developers.

[rkazants]

Hi @rghvsh, do you have update? Please ask questions if something is unclear.

Best regards,
Roman

[rghvsh]

Hi! @rkazants my exams got over today just will just brush up few things and start a pr tomorrow.

Sorry for the very long delay

Best regards,
Raghav

[p-wysocki]

No worries @rghvsh, we're just checking what's going on - we have to periodically check all tasks for activity. No reason to apologize! :)

[rkazants]

Hi @rghvsh, any update on this task?

[rghvsh]

Hi @rkazants, Yes i have started making the subgraphs.
Had a question the functions for case statement where will they be defined.

[rkazants]

Hi @rkazants, Yes i have started making the subgraphs. Had a question the functions for case statement where will they be defined.

Not sure about your question. Just create ov::Model for each branch and use nested decomposition of If operation from our opset. You can see extraction and creation of branch function for tf.raw_ops.If (a particular case of Case) in openvino\src\frontends\tensorflow\src\op\if.cpp

[rghvsh]

Hello @rkazants ! hope you're doing good

I wanted to run this by you is this the correct method to retrieve the desired branch type

    // retrieve body ov::Model for branches
    auto branches = node.get_attribute<std::string>("branches");
    // get desired branch index  
    auto cond = node.get_input(0)
    auto desired_branch_type = branches[cond]
    auto then_branch_body = translate_session->get_body_ov_model(desired_branch_type, ov_inputs);
    FRONT_END_GENERAL_CHECK(
        then_branch_body,
        "[TensorFlow Frontend] Internal error or incorrect input model. Cannot find branch body graph with name " +
            desired_branch_type);

[rkazants]

Hello @rkazants ! hope you're doing good

I wanted to run this by you is this the correct method to retrieve the desired branch type

    // retrieve body ov::Model for branches
    auto branches = node.get_attribute<std::string>("branches");
    // get desired branch index  
    auto cond = node.get_input(0)
    auto desired_branch_type = branches[cond]
    auto then_branch_body = translate_session->get_body_ov_model(desired_branch_type, ov_inputs);
    FRONT_END_GENERAL_CHECK(
        then_branch_body,
        "[TensorFlow Frontend] Internal error or incorrect input model. Cannot find branch body graph with name " +
            desired_branch_type);

Hi @rghvsh, you are correct. Keep this way. I anticipate your PR:)

Best regards,
Roman

[rghvsh]

Hey! @rkazants can we implement case like this via if operation we get the desired branch load it into then body and a non desired branch in else body and set condition as true?

    auto if_op = std::make_shared<If>(1);
    if_op->set_then_body(desired_branch_body);
    if_op->set_else_body(non_desired_branch_body);

Thanks
Raghav

[rkazants]

Hey! @rkazants can we implement case like this via if operation we get the desired branch load it into then body and a non desired branch in else body and set condition as true?

    auto if_op = std::make_shared<If>(1);
    if_op->set_then_body(desired_branch_body);
    if_op->set_else_body(non_desired_branch_body);

Thanks Raghav

Hi @rghvsh, in non_desired_branch_body you should create another if graph operation to check the next case value. It should nested If structure to represent tf.raw_ops.Case because we have no direct analogue for Case in OV opset.

Best regards,
Roman

[rghvsh]

Hey! @rkazants Case is a general implementation of If. I tried here to implement case via If operation.

In case we have desired branch index which needs to be executed. We can extract the desired branch from input and then setup an If operation and in then body we enter the desired branch and set condition as True so that always desired branch is executed.

For else branch we can have a placeholder which also is the desired branch but it will not be executed. I did this to cover edge cases such as the list of branches has length = 1 (just the desired branch) then also we can run it.

Please see whether the below implementation is correct or should I make some changes

    // retrieve body ov::Model for then and else branches
    auto branches = node.get_attribute<std::string>("branches");
    // get desired branch index  
    auto cond = node.get_input(0)
    auto desired_branch_type = branches[cond]
    auto placeholder_branch_type = branches[cond]

    auto desired_branch_body= translate_session->get_body_ov_model(desired_branch_type, ov_inputs);
            
    auto placeholder_branch_body = translate_session->get_body_ov_model(placeholder_branch_type, ov_inputs);

    // get condition output
    auto desired_params = desired_branch_body->get_parameters();
    auto placeholder_params = placeholder_branch_body->get_parameters();

    // create If operation and set body graphs
    auto if_op = std::make_shared<If>(1);
    if_op->set_then_body(desired_branch_body);
    if_op->set_else_body(placeholder_branch_body);

    // set inputs
    for (int ind = 0; ind < input_size; ++ind) {
        auto curr_input = node.get_input(ind + 1);
        auto desired_param = desired_params[ind];
        auto placeholder_param = else_params[ind];
        if_op->set_input(curr_input, desired_param, placeholder_param);
    }

    // set outputs
    auto desired_results = desired_branch_body->get_results();
    auto placeholder_results = placeholder_branch_body->get_results();
    int output_size = static_cast<int>(desired_results.size());
    for (int ind = 0; ind < output_size; ++ind) {
        if_op->set_output(desired_results[ind], placeholder_results[ind]);
    }

    auto ov_outputs = if_op->outputs();

    // set output tensor names
    for (size_t idx = 0; idx < ov_outputs.size(); ++idx) {
        ov_outputs[idx].get_tensor().set_names({node_name + ":" + std::to_string(idx)});
    }

    return ov_outputs;

Thanks
Raghav

[rkazants]

Hi @rghvsh, the idea is correct but you should also create else_branch_body (in your terms placeholder_branch_body ) ov::Model model from scratch in your translator where you put new If operation with then_branch from branches[ind + 1] and create else_brach_body from scratch again. It will be a sort of iteratively created nested-If strcuture while you iterate through all branches.

Please create PR and we can discuss it there during code review.

Best regards,
Roman

[rghvsh]

PR #22556

[shubdas9902]

.take

[github-actions]

Thank you for looking into this issue! Please let us know if you have any questions or require any help.

[shubdas9902]

Hi @rkazants,

I am new to open-source contributions and have been reviewing the documentation and guide for this issue. From what I understand, I need to create a file in the op directory, named case.cpp, where I will implement the support for the Case operation in the TensorFlow Frontend (TF FE).

Am I on the right track?

Thanks,
Shubham