while_op.cc

#include "caffe2/operators/while_op.h"

namespace caffe2 {

REGISTER_CPU_OPERATOR(While, WhileOp<CPUContext>);

OPERATOR_SCHEMA(While)
    .NumInputs(1, INT_MAX)
    .NumOutputs(0, INT_MAX)
    .SetDoc(R"DOC(
'While' control operator, first input is a scalar boolean blob that stores loop's
condition value. Accepts 'loop_net' (required) and 'cond_net' (optional) arguments for
loop's body and condition subnets respectively. If condition subnet is specified,
it is executed before the first and after each iteration. Subnets are executed in
the same workspace as 'While'.
    )DOC")
    .Arg("loop_net", "Net executed on each iteration")
    .Arg("cond_net", "Net to (re)compute condition value")
    .Input(0, "condition", "Scalar boolean condition")
    .AllowInplace([](int in, int out) -> bool { return true; });

} // namespace caffe2