Support Indexing for multiple return values from a Tango Step

[sunitharavi9]

🚀 The feature, motivation and pitch

If a Tango step returns multiple values, we currently have to pass all the returned values to another Step. The goal is to support indexing so that the user can select which returned values should be passed as an input to another Step.

Alternatives

No response

Additional context

No response

[epwalsh]

One alternative worth mentioning (which is currently possible) is to create an intermediate step which takes any kind of Sequence and just returns the item corresponding to a given index. So your config could look like:

sequence: {type: "sequence_step"},  # pretend this step returns a sequence of whatever
second_item: {type: "get_index", index: 1},  # this is the intermediate 'getter' step
final_step: {type: "whatever", input: {type: "ref", ref: second_item}},

[BigRedT]

I would like to second @isunitha98selvan's feature request for being able to index multiple returned values.

In the following example, InferenceStep needs to refer to individual outputs from ModelCreationStep. I know there is a workaround here - wrapping model and params into a single return object - but would be nice to have the flexibility!

@Step.register('init_model')
class ModelCreationStep(Step):
    DETERMINISTIC = False
    CACHEABLE = False

    def run(self,image_size):
        model = Model(image_size=image_size)
        params = get_params(model)
        return model, params

@Step.register('infer')
class InferenceStep(Step):
    DETERMINISTIC = False
    CACHEABLE = False

    def run(self,model,params,input):
        output = model.apply(params,input)
        return output

@epwalsh

[epwalsh]

Gotcha.. Hmm, in the JSON config, I'm thinking we'd support it this way:

{
  steps: {
    model_and_params: { type: "init_model" },
    predictions: { type: "infer", model: { type: "ref", ref: "model_and_params", key: 0 }, ... },
  }
}

The main thing here is the addition of a "key" field in the ref step object. The key could also be a string so we could pull objects out of dictionaries.

What do you think @BigRedT @AkshitaB @dirkgr?

[BigRedT]

That would work for me!

[dirkgr]

Just noticed this. This is a brilliant feature!