/*

* Licensed to the Apache Software Foundation (ASF) under one

* or more contributor license agreements. See the NOTICE file

* distributed with this work for additional information

* regarding copyright ownership. The ASF licenses this file

* to you under the Apache License, Version 2.0 (the

* "License"); you may not use this file except in compliance

* with the License. You may obtain a copy of the License at

*

* http://www.apache.org/licenses/LICENSE-2.0

*

* Unless required by applicable law or agreed to in writing,

* software distributed under the License is distributed on an

* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

* KIND, either express or implied. See the License for the

* specific language governing permissions and limitations

* under the License.

*/

#include <mxnet/io.h>

#include <mxnet/base.h>

#include <mxnet/ndarray.h>

#include <mxnet/operator.h>

#include <mxnet/operator_util.h>

#include <dmlc/logging.h>

#include <dmlc/optional.h>

#include <utility>

#include "./operator_common.h"

#include "./elemwise_op_common.h"

#include "../imperative/imperative_utils.h"

#include "./subgraph_op_common.h"

namespace mxnet {

namespace op {

struct ForeachParam : public dmlc::Parameter<ForeachParam> {

int num_args;

int num_outputs;

int num_out_data;

// The location of states in the subgraph inputs.

mxnet::Tuple<dim_t> in_state_locs;

// The location of data arrays in the subgraph inputs.

mxnet::Tuple<dim_t> in_data_locs;

// The location of remaining arrays in the subgraph inputs.

mxnet::Tuple<dim_t> remain_locs;

DMLC_DECLARE_PARAMETER(ForeachParam) {

DMLC_DECLARE_FIELD(num_args).set_lower_bound(1).describe("Number of inputs.");

DMLC_DECLARE_FIELD(num_outputs).describe("The number of outputs of the subgraph.");

DMLC_DECLARE_FIELD(num_out_data).describe("The number of output data of the subgraph.");

DMLC_DECLARE_FIELD(in_state_locs).describe("The locations of loop states among the inputs.");

DMLC_DECLARE_FIELD(in_data_locs).describe("The locations of input data among the inputs.");

DMLC_DECLARE_FIELD(remain_locs).describe("The locations of remaining data among the inputs.");

}

}; // struct ForeachParam

DMLC_REGISTER_PARAMETER(ForeachParam);

class ForeachState : public LoopState {

public:

ForeachParam params;

int num_iterations;

ForeachState(const nnvm::Symbol& g, const ForeachParam& params) : LoopState(g, false) {

this->params = params;

}

};

static void ForeachComputeExCPU(const OpStatePtr& state_ptr,

const OpContext& ctx,

const std::vector<NDArray>& inputs,

const std::vector<OpReqType>& req,

const std::vector<NDArray>& outputs) {

ForeachState& state = state_ptr.get_state<ForeachState>();

const ForeachParam& params = state.params;

const size_t iter_dim = 0;

CHECK_EQ(outputs.size(), (size_t)params.num_outputs);

CHECK_GT(params.in_data_locs.ndim(), 0);

size_t len = inputs[0].shape()[iter_dim];

state.num_iterations = len;

for (int i = 1; i < params.in_data_locs.ndim(); i++)

CHECK_EQ(inputs[i].shape()[iter_dim], len);

for (size_t i = 0; i < (size_t)params.num_out_data; i++)

CHECK_EQ(len, outputs[i].shape()[iter_dim]);

for (const auto& arr : outputs)

CHECK_EQ(arr.storage_type(), kDefaultStorage)

<< "The for operator doesn't support the sparse format";

// Initialize the outputs of the subgraph is a little trickier.

// The states from the previous iteration are used as the inputs of the next

// iteration, so I have to maintain two arrays, so the inputs and outputs

// of the subgraph share the same memory.

std::vector<NDArray> subg_outputs1(outputs.size());

std::vector<NDArray> subg_outputs2(outputs.size());

std::vector<NDArray>* subg_outputs[2]{&subg_outputs1, &subg_outputs2};

// If the length is an odd number, the last iteration will use the first set

// of outputs. In this way, we don't need to copy the results from the

// subgraph to the final outputs of the loop.

if (len % 2 == 1) {

for (size_t i = params.num_out_data; i < subg_outputs1.size(); i++) {

subg_outputs1[i] = outputs[i];

subg_outputs2[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), true, outputs[i].dtype());

}

} else {

// Otherwise, we'll use the second set of outputs.

for (size_t i = params.num_out_data; i < subg_outputs1.size(); i++) {

subg_outputs1[i] = NDArray(outputs[i].shape(), outputs[i].ctx(), true, outputs[i].dtype());

subg_outputs2[i] = outputs[i];

}

}

// Initialize the inputs for the subgraph.

// In each iteration, we need to update the subgraph inputs for input data

// and the loop states.

std::vector<NDArray> subg_inputs(inputs.size());

// The remaining arrays (other than input data and states) only need to be set once.

for (int j = 0; j < params.remain_locs.ndim(); j++) {

CHECK_LT(params.remain_locs[j], subg_inputs.size());

subg_inputs[params.remain_locs[j]] =

inputs[j + params.in_data_locs.ndim() + params.in_state_locs.ndim()];

}

// Here we iterate over the first dimension of the first input array.

for (size_t i = 0; i < len; i++) {

// Initialize outputs for the subgraph.

std::vector<NDArray>* subg_out_curr = subg_outputs[i % 2];

std::vector<NDArray>* subg_out_prev = subg_outputs[(i + 1) % 2];

for (int j = 0; j < params.num_out_data; j++)

(*subg_out_curr)[j] = outputs[j].At(i);

// When recording for backward computation, we should make sure

// that output arrays are actually different in each iteration.

if (ctx.need_grad && i < len - 1) {

for (size_t j = params.num_out_data; j < subg_out_curr->size(); j++)

(*subg_out_curr)[j] =

NDArray(outputs[j].shape(), outputs[j].ctx(), true, outputs[j].dtype());

} else if (ctx.need_grad && i == len - 1) {

// For the last iteration, we need to write data to the output array

// directly.

for (size_t j = params.num_out_data; j < subg_out_curr->size(); j++)

(*subg_out_curr)[j] = outputs[j];

}

// Initialize inputs for the subgraph.

// Get a slice from the input data arrays.

for (int j = 0; j < params.in_data_locs.ndim(); j++) {

size_t loc = params.in_data_locs[j];

subg_inputs[loc] = inputs[j].At(i);

}

// For the rest of the iterations, the states are the outputs

// from the previous iteration.

if (i > 0) {

for (size_t j = params.num_out_data; j < subg_out_prev->size(); j++) {

size_t idx = j - params.num_out_data;

CHECK_LT(params.in_state_locs[idx], subg_inputs.size());

subg_inputs[params.in_state_locs[idx]] = (*subg_out_prev)[j];

}

} else {

for (int j = 0; j < params.in_state_locs.ndim(); j++) {

CHECK_LT(params.in_state_locs[j], subg_inputs.size());

subg_inputs[params.in_state_locs[j]] = inputs[j + params.in_data_locs.ndim()];

}

}

state.Forward(i, subg_inputs, req, *subg_out_curr, ctx.need_grad);

}

}

static void ForeachGradComputeExCPU(const OpStatePtr& state_ptr,

const OpContext& ctx,

const std::vector<NDArray>& inputs,

const std::vector<OpReqType>& req,

const std::vector<NDArray>& outputs) {

ForeachState& state = state_ptr.get_state<ForeachState>();

const ForeachParam& params = state.params;

CHECK_EQ(outputs.size(), (size_t)params.num_args - 1);

CHECK_GT(params.in_data_locs.ndim(), 0);

for (const auto& arr : outputs)

CHECK_EQ(arr.storage_type(), kDefaultStorage)

<< "The for operator doesn't support the sparse format";

int len = state.num_iterations;

size_t num_output_data = params.num_out_data;

// In backward computation, we need to run iterations from backwards.

std::vector<NDArray> subg_ograds(params.num_outputs);

std::vector<NDArray> subg_igrads(outputs.size());

for (size_t i = num_output_data; i < subg_ograds.size(); i++)

subg_ograds[i] = inputs[i];

std::vector<OpReqType> subg_req(req.size());

for (auto r : req)

CHECK_NE(r, kWriteInplace);

// There are three types of arrays in igrads.

// * data gradients.

// * loop variable gradients.

// * remaining variable gradients.

// They are in the following order:

// [data vars], [loop vars], [remaining vars]

// [remaining vars]

for (int i = 0; i < params.remain_locs.ndim(); i++) {

size_t loc = params.remain_locs[i];

size_t orig_loc = i + params.in_data_locs.ndim() + params.in_state_locs.ndim();

subg_igrads[loc] = outputs[orig_loc];

subg_req[loc] = req[orig_loc];

}

for (int iter_num = len - 1; iter_num >= 0; iter_num--) {

for (int i = 0; i < params.num_out_data; i++)

subg_ograds[i] = inputs[i].At(iter_num);

if (iter_num < len - 1) {

// For the rest of the iterations, we should add graidents to the

// remaining vars.

for (int i = 0; i < params.remain_locs.ndim(); i++) {

size_t loc = params.remain_locs[i];

subg_req[loc] = kAddTo;

}

}

// [data vars]

for (int i = 0; i < params.in_data_locs.ndim(); i++) {

size_t loc = params.in_data_locs[i];

subg_igrads[loc] = outputs[i].At(iter_num);

subg_req[loc] = req[i];

}

// [loop vars]

for (int i = 0; i < params.in_state_locs.ndim(); i++) {

size_t loc = params.in_state_locs[i];

const NDArray& output = outputs[i + params.in_data_locs.ndim()];

if (iter_num != 0) {

// For state gradients, we need to allocate new NDArrays

// because intermediate state gradients won't be returned to the users.

subg_igrads[loc] = NDArray(output.shape(), output.ctx(), true, output.dtype());

} else {

subg_igrads[loc] = output;

}

// For the first iteration, we need to use the request provided by

// the user to write state gradients to the outputs.

subg_req[loc] = iter_num != 0 ? kWriteTo : req[i + params.in_data_locs.ndim()];

}

state.Backward(iter_num, subg_ograds, subg_req, subg_igrads);

size_t num_states = subg_ograds.size() - num_output_data;

for (size_t i = 0; i < num_states; i++) {

size_t loc = params.in_state_locs[i];

CHECK_LT(loc, subg_igrads.size());

subg_ograds[i + num_output_data] = subg_igrads[loc];

}

}

state.Cleanup();

}

template <typename T>

static void remap(const std::vector<T>& op_in,

size_t start,

const mxnet::Tuple<dim_t>& locs,

std::vector<T>* subg_in) {

auto op_in_it = op_in.begin() + start;

for (int i = 0; i < locs.ndim(); i++) {

dim_t loc = locs[i];

subg_in->at(loc) = *(op_in_it + i);

}

}

static inline mxnet::TShape SliceFirstDim(const mxnet::TShape& s) {

if (s.ndim() > 1) {

return mxnet::TShape(s.begin() + 1, s.end());

} else {

return mxnet::TShape(mshadow::Shape1(1));

}

}

static bool ForeachShape(const nnvm::NodeAttrs& attrs,

mxnet::ShapeVector* in_shape,

mxnet::ShapeVector* out_shape) {

const ForeachParam& params = nnvm::get<ForeachParam>(attrs.parsed);

CHECK_EQ(out_shape->size(), (size_t)params.num_outputs);

CHECK_EQ(attrs.subgraphs.size(), 1U);

mxnet::ShapeVector subg_in_shape(in_shape->size());

// data shape

std::vector<bool> data_1d(params.in_data_locs.ndim(), false);

for (int i = 0; i < params.in_data_locs.ndim(); i++) {

size_t loc = params.in_data_locs[i];

if (in_shape->at(i).ndim() == 1)

data_1d[i] = true;

subg_in_shape[loc] = SliceFirstDim(in_shape->at(i));

}

// state shape

remap(*in_shape, params.in_data_locs.ndim(), params.in_state_locs, &subg_in_shape);

// remaining shape

remap(*in_shape,

params.in_data_locs.ndim() + params.in_state_locs.ndim(),

params.remain_locs,

&subg_in_shape);

mxnet::ShapeVector subg_out_shape = *out_shape;

for (int i = 0; i < params.num_out_data; i++) {

mxnet::TShape shape = subg_out_shape[i];

// If we don't have shape info, we don't need to do anything.

if (!mxnet::ndim_is_known(shape))

continue;

subg_out_shape[i] = SliceFirstDim(shape);

}

bool infer_success = InferSubgraphShape(*attrs.subgraphs[0], &subg_in_shape, &subg_out_shape);

// After inference, we need to move inferred information back to in_shape and

// out_shape.

// For the shape of output data.

size_t len = in_shape->at(0)[0];

for (int i = 0; i < params.num_out_data; i++) {

// If the output shape isn't inferred, we don't need to propogate the info.

const auto& g_out_shape = subg_out_shape[i];

if (!mxnet::ndim_is_known(g_out_shape))

continue;

auto out = mxnet::TShape(g_out_shape.ndim() + 1, -1);

out[0] = len;

for (int i = 1; i < out.ndim(); i++)

out[i] = g_out_shape[i - 1];

SHAPE_ASSIGN_CHECK(*out_shape, i, out);

}

// For the shape of output states.

for (size_t i = params.num_out_data; i < subg_out_shape.size(); i++)

SHAPE_ASSIGN_CHECK(*out_shape, i, subg_out_shape[i]);

// For the shape of input data.

for (int i = 0; i < params.in_data_locs.ndim(); i++) {

size_t loc = params.in_data_locs[i];

const auto& shape = subg_in_shape[loc];

// If the input data shape isn't inferred, we don't need to propogate the

// info.

if (!mxnet::ndim_is_known(shape))

continue;

if (data_1d[i]) {

mxnet::TShape s(1, -1);

s[0] = len;

SHAPE_ASSIGN_CHECK(*in_shape, i, s);

} else {

auto in = mxnet::TShape(shape.ndim() + 1, -1);

in[0] = len;

for (int i = 1; i < in.ndim(); i++)

in[i] = shape[i - 1];

SHAPE_ASSIGN_CHECK(*in_shape, i, in);

}

}

// For the shape of state.

for (int i = 0; i < params.in_state_locs.ndim(); i++) {

size_t loc = params.in_state_locs[i];

SHAPE_ASSIGN_CHECK(*in_shape, i + params.in_data_locs.ndim(), subg_in_shape[loc]);

}

// For the shape of remaining data.

for (int i = 0; i < params.remain_locs.ndim(); i++) {

size_t loc = params.remain_locs[i];

SHAPE_ASSIGN_CHECK(*in_shape,

i + params.in_data_locs.ndim() + params.in_state_locs.ndim(),

subg_in_shape[loc]);

}

if (infer_success) {

size_t num_states = out_shape->size() - params.num_out_data;

for (size_t i = 0; i < num_states; i++) {

CHECK_EQ((*out_shape)[i + params.num_out_data], (*in_shape)[i + params.in_data_locs.ndim()]);

}

}

// Check if we have inferred the shapes correctly.

return infer_success;

}

static bool ForeachType(const nnvm::NodeAttrs& attrs,

std::vector<int>* in_type,

std::vector<int>* out_type) {

const ForeachParam& params = nnvm::get<ForeachParam>(attrs.parsed);

CHECK_EQ(out_type->size(), (size_t)params.num_outputs);

CHECK_EQ(attrs.subgraphs.size(), 1U);

std::vector<int> subg_in_type(in_type->size(), 0);

remap(*in_type, 0, params.in_data_locs, &subg_in_type);

remap(*in_type, params.in_data_locs.ndim(), params.in_state_locs, &subg_in_type);

remap(*in_type,

params.in_data_locs.ndim() + params.in_state_locs.ndim(),

params.remain_locs,

&subg_in_type);

bool success = InferSubgraphDataType(*attrs.subgraphs[0], &subg_in_type, out_type);

for (int i = 0; i < params.in_data_locs.ndim(); i++) {

size_t loc = params.in_data_locs[i];

TYPE_ASSIGN_CHECK(*in_type, i, subg_in_type[loc]);

}

for (int i = 0; i < params.in_state_locs.ndim(); i++) {

size_t loc = params.in_state_locs[i];

TYPE_ASSIGN_CHECK(*in_type, i + params.in_data_locs.ndim(), subg_in_type[loc]);

}

for (int i = 0; i < params.remain_locs.ndim(); i++) {

size_t loc = params.remain_locs[i];

TYPE_ASSIGN_CHECK(

*in_type, i + params.in_data_locs.ndim() + params.in_state_locs.ndim(), subg_in_type[loc]);

}

return success;

}

static bool ForeachStorageType(const nnvm::NodeAttrs& attrs,

const int dev_mask,

DispatchMode* dispatch_mode,

std::vector<int>* in_attrs,

std::vector<int>* out_attrs) {

const ForeachParam& params = nnvm::get<ForeachParam>(attrs.parsed);

CHECK_EQ(out_attrs->size(), (size_t)params.num_outputs);

CHECK_EQ(attrs.subgraphs.size(), 1U);

std::vector<int> subg_in_attrs(in_attrs->size(), kUndefinedStorage);

remap(*in_attrs, 0, params.in_data_locs, &subg_in_attrs);

remap(*in_attrs, params.in_data_locs.ndim(), params.in_state_locs, &subg_in_attrs);

remap(*in_attrs,

params.in_data_locs.ndim() + params.in_state_locs.ndim(),

params.remain_locs,

&subg_in_attrs);

bool success =

InferSubgraphStorage(*attrs.subgraphs[0], dev_mask, dispatch_mode, &subg_in_attrs, out_attrs);

for (int i = 0; i < params.in_data_locs.ndim(); i++) {

size_t loc = params.in_data_locs[i];

STORAGE_TYPE_ASSIGN_CHECK(*in_attrs, i, subg_in_attrs[loc]);

}

for (int i = 0; i < params.in_state_locs.ndim(); i++) {

size_t loc = params.in_state_locs[i];

STORAGE_TYPE_ASSIGN_CHECK(*in_attrs, i + params.in_data_locs.ndim(), subg_in_attrs[loc]);

}

for (int i = 0; i < params.remain_locs.ndim(); i++) {

size_t loc = params.remain_locs[i];

STORAGE_TYPE_ASSIGN_CHECK(*in_attrs,

i + params.in_data_locs.ndim() + params.in_state_locs.ndim(),

subg_in_attrs[loc]);

}

return success;

}

static bool BackwardForeachStorageType(const nnvm::NodeAttrs& attrs,

const int dev_mask,

DispatchMode* dispatch_mode,

std::vector<int>* in_attrs,

std::vector<int>* out_attrs) {

const ForeachParam& params = nnvm::get<ForeachParam>(attrs.parsed);

CHECK_EQ(out_attrs->size(), (size_t)params.num_args - 1);

CHECK_EQ(in_attrs->size(), (size_t)params.num_args - 1 + params.num_outputs * 2);

CHECK_EQ(attrs.subgraphs.size(), 1U);

CachedOp op(*attrs.subgraphs[0], std::vector<std::pair<std::string, std::string> >());

// map the operator inputs to the subgraph inputs.

std::vector<int> subg_forward_ins(params.num_args - 1, kUndefinedStorage);

remap(*in_attrs, params.num_outputs, params.in_data_locs, &subg_forward_ins);

remap(*in_attrs,

params.num_outputs + params.in_data_locs.ndim(),

params.in_state_locs,

&subg_forward_ins);

remap(*in_attrs,

params.num_outputs + params.in_data_locs.ndim() + params.in_state_locs.ndim(),

params.remain_locs,

&subg_forward_ins);

// Copy backward input storage to backward subgraph input storage.

std::vector<int> subg_in_attrs = *in_attrs;

for (size_t i = 0; i < subg_forward_ins.size(); i++)

subg_in_attrs[i + params.num_outputs] = subg_forward_ins[i];

return op.BackwardStorageType(attrs, dev_mask, dispatch_mode, &subg_in_attrs, out_attrs);

}

static OpStatePtr CreateForeachState(const NodeAttrs& attrs,

Context ctx,

const mxnet::ShapeVector& ishape,

const std::vector<int>& itype) {

const ForeachParam& params = nnvm::get<ForeachParam>(attrs.parsed);

return OpStatePtr::Create<ForeachState>(*attrs.subgraphs[0], params);

}

static std::vector<nnvm::NodeEntry> ForeachGradient(const nnvm::ObjectPtr& n,

const std::vector<nnvm::NodeEntry>& ograds) {

ElemwiseGradUseInOut fgrad{"_backward_foreach"};

std::vector<nnvm::NodeEntry> entries = fgrad(n, ograds);

entries[0].node->attrs.subgraphs = n->attrs.subgraphs;

return entries;

}

struct WhileLoopParam : public dmlc::Parameter<WhileLoopParam> {

int num_args;

int num_outputs;

int num_out_data;

int max_iterations;

// `cond' and `func' each takes a subset of while_loop's inputs as that to their subgraphs

// `cond_input_locs' contains indices of inputs fed to `cond', and

// `func_input_locs' contains indices of inputs fed to `func'.

// `func_var_locs' are indices in which input "variables" are stored in func's inputs.

mxnet::Tuple<dim_t> cond_input_locs;

mxnet::Tuple<dim_t> func_input_locs;

mxnet::Tuple<dim_t> func_var_locs;

DMLC_DECLARE_PARAMETER(WhileLoopParam) {

DMLC_DECLARE_FIELD(num_args).set_lower_bound(2).describe(

"Number of input arguments, including cond and func as two symbol inputs.");

DMLC_DECLARE_FIELD(num_outputs)

.set_lower_bound(1)

.describe("The number of outputs of the subgraph.");

DMLC_DECLARE_FIELD(num_out_data)

.set_lower_bound(0)

.describe("The number of outputs from the function body.");

DMLC_DECLARE_FIELD(max_iterations).set_lower_bound(1).describe("Maximum number of iterations.");

DMLC_DECLARE_FIELD(cond_input_locs)

.describe("The locations of cond's inputs in the given inputs.");

DMLC_DECLARE_FIELD(func_input_locs)

.describe("The locations of func's inputs in the given inputs.");

DMLC_DECLARE_FIELD(func_var_locs).describe("The locations of loop_vars among func's inputs.");

}

template <typename T>

bool sync_in_out(std::vector<T>* in,

std::vector<T>* out,

std::function<bool(const T&)> is_empty) const {

for (int i = this->num_out_data; i < this->num_outputs; ++i) {

// each out->at(i) is a params, loop_var

T& x = in->at(this->func_input_locs[this->func_var_locs[i - this->num_out_data]]);

T& y = out->at(i);

fill_value(&x, &y, is_empty(x), is_empty(y));

}

return true;

}

}; // struct WhileLoopParam

DMLC_REGISTER_PARAMETER(WhileLoopParam);

class WhileLoopState : public LoopState {

public:

WhileLoopParam params;

size_t n_iterations; // the actual number of steps taken in this while loop, <= max_iterations

CachedOpPtr cond_op;

// abbrev for output_input_mapping

// indicates to which index the output of `func' will be copied to the input of `cond'

std::vector<int> oi_map;

WhileLoopState(const WhileLoopParam& params, const nnvm::Symbol& cond, const nnvm::Symbol& func)

: LoopState(func),

params(params),

n_iterations(0U),

cond_op(LoopState::MakeSharedOp(cond)),

oi_map(params.func_var_locs.ndim(), -1) {

const mxnet::Tuple<dim_t>& func_input_locs = params.func_input_locs;

const mxnet::Tuple<dim_t>& func_var_locs = params.func_var_locs;

const mxnet::Tuple<dim_t>& cond_input_locs = params.cond_input_locs;

for (int i = 0; i < func_var_locs.ndim(); ++i) {

dim_t pos_i = func_input_locs[func_var_locs[i]];

for (int j = 0; j < cond_input_locs.ndim(); ++j) {

dim_t pos_j = cond_input_locs[j];

if (pos_i == pos_j) {

this->oi_map[i] = j;

}

}

}

}

};

static void WhileLoopComputeExCPU(const OpStatePtr& state_ptr,

const OpContext& ctx,

const std::vector<NDArray>& inputs,

const std::vector<OpReqType>& req,

const std::vector<NDArray>& outputs) {

// The argument `inputs' are loop_vars and other inputs

// loop_vars are stored in stored in `loop_vars_locs'

// The argument `outputs' are output and new_loop_vars

// [0: num_out_data) are outputs at each step.

// [num_out_data: ) are new_loop_vars

// TODO(Junru): avoid dynamic NDArray allocation

WhileLoopState& state = state_ptr.get_state<WhileLoopState>();

const WhileLoopParam& params = state.params;

// a helper function, converting std::vector<NDArray> to std::vector<NDArray*>

const auto to_ptr_vec = [](std::vector<NDArray>& in, std::vector<NDArray*>* out) {

out->clear();

out->reserve(in.size());

std::transform(

std::begin(in), std::end(in), std::back_inserter(*out), [](NDArray& a) { return &a; });

};

// sanity checks

CHECK_EQ(inputs.size() + 2U, (size_t)params.num_args);

CHECK_EQ(outputs.size(), (size_t)params.num_outputs);

CHECK_EQ(outputs.size(), req.size());

// construct inputs and outputs for cond

std::vector<NDArray> cond_inputs, cond_outputs = {NDArray()};

extract_by_loc(inputs, params.cond_input_locs, &cond_inputs);

std::vector<NDArray*> cond_input_ptr, cond_output_ptr;

to_ptr_vec(cond_inputs, &cond_input_ptr);

to_ptr_vec(cond_outputs, &cond_output_ptr);

// construct inputs and outputs for func

std::vector<NDArray> func_inputs, func_outputs(outputs.size());

extract_by_loc(inputs, params.func_input_locs, &func_inputs);

for (size_t& step = state.n_iterations = 0; step < (size_t)params.max_iterations; ++step) {

CHECK(inputs.size() > 0) << "while loop forward requires at least 1 input";

Context default_ctx = inputs[0].ctx();

state.cond_op->Forward(nullptr, cond_input_ptr, cond_output_ptr, default_ctx);

if (!as_bool_scalar(*cond_output_ptr[0])) {

break;

}

// we create func_outputs for the current step:

for (size_t i = 0; i < outputs.size(); ++i) {

func_outputs[i] = NDArray(outputs[i].ctx(), outputs[i].dtype());

}

state.Forward(step, func_inputs, req, func_outputs, ctx.need_grad);

if (step == 0) {

for (int i = 0; i < params.num_out_data; ++i) {

func_outputs[i].WaitToRead();

if (!shape_is_known(func_outputs[i].shape())) {

func_outputs[i].SetShapeFromChunk();

}

mxnet::TShape step_shape = func_outputs[i].shape();

mxnet::TShape shape(step_shape.ndim() + 1, 0);

shape[0] = params.max_iterations;

for (int j = 0; j < step_shape.ndim(); ++j) {

shape[j + 1] = step_shape[j];

}

const_cast<NDArray&>(outputs[i]).Init(shape);

}

}

for (int i = 0; i < params.num_out_data; ++i) {

NDArray first_slot = outputs[i].At(step);

mxnet::CopyFromTo(func_outputs[i], &first_slot);

}

// func_inputs on the next step:

// the output (new_loop_vars) will become the new inputs (loop_vars)

for (size_t i = params.num_out_data; i < outputs.size(); ++i) {

int j = params.func_var_locs[i - params.num_out_data];

func_inputs[j] = func_outputs[i];

int k = state.oi_map[i - params.num_out_data];

if (k != -1) {

// I actually don't need to update cond_inputs

cond_inputs[k] = func_outputs[i];

cond_input_ptr[k] = &func_outputs[i];

}

}

}

// copy output data to `outputs'

// case 1: at least one step is executed,

// the final_loop_vars must be stored in func_inputs

// case 2: no step is executed

// the final_loop_vars is the same as loop_vars, which are also stored in func_inputs

// therefore, we copy func_inputs[:] to outputs[num_out_data: ]

for (size_t i = params.num_out_data; i < outputs.size(); ++i) {

size_t j = params.func_var_locs[i - params.num_out_data];

if (!shape_is_known(outputs[i].shape())) {

const_cast<NDArray&>(outputs[i]).Init(func_inputs[j].shape());

}

mxnet::CopyFromTo(func_inputs[j], &outputs[i]);

}

for (int i = 0; i < params.num_out_data; ++i) {

const_cast<NDArray&>(outputs[i]).SetShapeFromChunk();

}

if (state.n_iterations == 0) {

for (const auto& output : outputs) {

if (!shape_is_known(output.shape())) {

const_cast<NDArray&>(output).ReshapeAndAlloc({1});

}

}

}

}

static void WhileLoopGradComputeExCPU(const OpStatePtr& state_ptr,

const OpContext& ctx,

const std::vector<NDArray>& inputs,

const std::vector<OpReqType>& _req,

const std::vector<NDArray>& _outputs) {

// inputs are dl / df(x)

// outputs are dl / dx

// where f is the current function,

// x is the input to the current function,

// TODO(Junru): avoid dynamic NDArray allocation

WhileLoopState& state = state_ptr.get_state<WhileLoopState>();

const WhileLoopParam& params = state.params;

// sanity checks

CHECK_EQ(_outputs.size() + 2U, (size_t)params.num_args);

CHECK_EQ(_outputs.size(), _req.size());

for (auto x : _req) {

CHECK_NE(x, kWriteInplace);

}

std::vector<NDArray> outputs;

std::vector<OpReqType> req;

extract_by_loc(_outputs, params.func_input_locs, &outputs);

extract_by_loc(_req, params.func_input_locs, &req);

if (state.n_iterations == 0) {

for (int i = params.num_out_data; i < params.num_outputs; ++i) {

int j = params.func_var_locs[i - params.num_out_data];

mxnet::CopyFromTo(inputs[i], &outputs[j]);

}

state.Cleanup();

return;

}

// collect var_locs and out_locs, positions other than var_locs are out_locs, i.e.

// [0, var_locs[0])

// (var_locs[1], var_locs[2]),

// (var_locs[2], var_locs[3]),

// ...

// (var_locs[-2], var_locs[-1] = params.num_args - 2)

std::vector<dim_t> var_locs(params.func_var_locs.begin(), params.func_var_locs.end());

var_locs.push_back((dim_t)params.num_args - 2U);

sort(var_locs.begin(), var_locs.end());

// vectors for the backward loop

std::vector<NDArray> ograds(params.num_outputs);

std::vector<NDArray> igrads(outputs.size());

std::vector<OpReqType> iter_req(req.size());

for (int i = params.num_out_data; i < params.num_outputs; ++i)

ograds[i] = inputs[i];

const int n_iter = state.n_iterations;

for (int step = n_iter - 1; step >= 0; --step) {

// ograds[ : num_out_data] = inputs[ : num_out_data][step]

// ograds[num_out_data: ] is maintained in the end of each loop

std::transform(std::begin(inputs),

std::begin(inputs) + params.num_out_data,

std::begin(ograds),

[step](const NDArray& a) { return a.At(step); });

// igrads[i] =

// outputs[i] (step == 0)

// outputs[i] (step != 0 && i not in loop_var_locs)

// ArrayLike(outputs[i]) (step != 0 && i in loop_var_locs)

// iter_req =

// kWriteTo (step != 0 && i in loop_var_locs)

// req[i] (step == 0 && i in loop_var_locs)

// kAddTo (step != n_iters - 1 && i not in loop_var_locs)

// req[i] (step == n_iters - 1 && i not in loop_var_locs)

{

size_t i = 0;

for (size_t loc : var_locs) {

for (; i < loc; ++i) {

// locs other that var_locs

igrads[i] = outputs[i];

iter_req[i] = (step + 1 == n_iter || req[i] == kNullOp) ? req[i] : kAddTo;

}

if (i < (size_t)params.num_args - 2U) {

// a var

igrads[i] = (step == 0) ?

outputs[i] :

NDArray(outputs[i].shape(), outputs[i].ctx(), true, outputs[i].dtype());

iter_req[i] = (step == 0 || req[i] == kNullOp) ? req[i] : kWriteTo;

++i;

} else {

break;

}

}

}

state.Backward(step, ograds, iter_req, igrads);

for (int i = params.num_out_data; i < params.num_outputs; ++i) {

size_t j = params.func_var_locs[i - params.num_out_data];

ograds[i] = igrads[j];

}

}

state.Cleanup();

}

static bool WhileLoopType(const nnvm::NodeAttrs& attrs,

std::vector<int>* in_type,

std::vector<int>* out_type) {

const WhileLoopParam& params = nnvm::get<WhileLoopParam>(attrs.parsed);

static const std::function<bool(const int&)> is_udf = is_type_udf;

CHECK_EQ(in_type->size() + 2U, (size_t)params.num_args);

CHECK_EQ(out_type->size(), (size_t)params.num_outputs);

CHECK_EQ(attrs.subgraphs.size(), 2U);

CHECK_EQ(attrs.subgraphs[0]->outputs.size(), 1U);

std::vector<int> cond_in_type;

std::vector<int> func_in_type;

extract_by_loc(*in_type, params.cond_input_locs, &cond_in_type);

extract_by_loc(*in_type, params.func_input_locs, &func_in_type);

std::vector<int> cond_out_type = {-1};

CHECK(params.sync_in_out(in_type, out_type, is_udf));

bool succ_0 = InferSubgraphDataType(*attrs.subgraphs[0], &cond_in_type, &cond_out_type);

CHECK(params.sync_in_out(in_type, out_type, is_udf));

CHECK(sync_in_in(params.cond_input_locs, in_type, &cond_in_type, is_udf));

bool succ_1 = InferSubgraphDataType(*attrs.subgraphs[1], &func_in_type, out_type);

CHECK(params.sync_in_out(in_type, out_type, is_udf));

CHECK(sync_in_in(params.func_input_locs, in_type, &func_in_type, is_udf));

return succ_0 && succ_1;

}

static bool WhileLoopStorageType(const nnvm::NodeAttrs& attrs,

const int dev_mask,

DispatchMode* dispatch_mode,

std::vector<int>* in_attrs,

std::vector<int>* out_attrs) {

const WhileLoopParam& params = nnvm::get<WhileLoopParam>(attrs.parsed);

static const std::function<bool(const int&)> is_udf = is_stype_udf;

CHECK_EQ(in_attrs->size() + 2U, (size_t)params.num_args);

CHECK_EQ(out_attrs->size(), (size_t)params.num_outputs);

CHECK_EQ(attrs.subgraphs.size(), 2U);

CHECK_EQ(attrs.subgraphs[0]->outputs.size(), 1U);

std::vector<int> cond_in_attrs;

std::vector<int> func_in_attrs;

extract_by_loc(*in_attrs, params.cond_input_locs, &cond_in_attrs);

extract_by_loc(*in_attrs, params.func_input_locs, &func_in_attrs);

std::vector<int> cond_out_attrs = {kDefaultStorage};

DispatchMode cond_mode = DispatchMode::kUndefined;

DispatchMode func_mode = DispatchMode::kUndefined;

*dispatch_mode = DispatchMode::kFComputeEx;

CHECK(params.sync_in_out(in_attrs, out_attrs, is_udf));

bool succ_0 = InferSubgraphStorage(

*attrs.subgraphs[0], dev_mask, &cond_mode, &cond_in_attrs, &cond_out_attrs);

CHECK(params.sync_in_out(in_attrs, out_attrs, is_udf));

CHECK(sync_in_in(params.cond_input_locs, in_attrs, &cond_in_attrs, is_udf));

bool succ_1 =

InferSubgraphStorage(*attrs.subgraphs[1], dev_mask, &func_mode, &func_in_attrs, out_attrs);

CHECK(params.sync_in_out(in_attrs, out_attrs, is_udf));

CHECK(sync_in_in(params.func_input_locs, in_attrs, &func_in_attrs, is_udf));

return succ_0 && succ_1;

}

static bool BackwardWhileLoopStorageType(const nnvm::NodeAttrs& attrs,

const int dev_mask,

DispatchMode* dispatch_mode,

std::vector<int>* in_attrs,

std::vector<int>* out_attrs) {

// `cond' is not backwarded, don't check

const WhileLoopParam& params = nnvm::get<WhileLoopParam>(attrs.parsed);

CHECK_EQ(out_attrs->size() + 2U, (size_t)params.num_args);

CHECK_EQ(attrs.subgraphs.size(), 2U);

CachedOp op(*attrs.subgraphs[1], {});

return op.BackwardStorageType(attrs, dev_mask, dispatch_mode, in_attrs, out_attrs);

}

static OpStatePtr CreateWhileLoopState(const NodeAttrs& attrs,

Context ctx,

const mxnet::ShapeVector& ishape,

const std::vector<int>& itype) {

const WhileLoopParam& params = nnvm::get<WhileLoopParam>(attrs.parsed);

return OpStatePtr::Create<WhileLoopState>(params, *attrs.subgraphs[0], *attrs.subgraphs[1]);

}

static std::vector<nnvm::NodeEntry> WhileLoopGradient(const nnvm::ObjectPtr& n,

const std::vector<nnvm::NodeEntry>& ograds) {

ElemwiseGradUseInOut fgrad{"_backward_while_loop"};

std::vector<nnvm::NodeEntry> entries = fgrad(n, ograds);

entries[0].node->attrs.subgraphs = n->attrs.subgraphs;

return entries;

}

struct CondParam : public dmlc::Parameter<CondParam> {

int num_args;

int num_outputs;

mxnet::Tuple<dim_t> cond_input_locs;

mxnet::Tuple<dim_t> then_input_locs;

mxnet::Tuple<dim_t> else_input_locs;

DMLC_DECLARE_PARAMETER(CondParam) {

DMLC_DECLARE_FIELD(num_args).set_lower_bound(3).describe(

"Number of input arguments, including cond, then and else as three symbol inputs.");

DMLC_DECLARE_FIELD(num_outputs)

.set_lower_bound(1)

.describe("The number of outputs of the subgraph.");

DMLC_DECLARE_FIELD(cond_input_locs)

.describe("The locations of cond's inputs in the given inputs.");

DMLC_DECLARE_FIELD(then_input_locs)

.describe("The locations of then's inputs in the given inputs.");

DMLC_DECLARE_FIELD(else_input_locs)

.describe("The locations of else's inputs in the given inputs.");

}

}; // struct CondParam

DMLC_REGISTER_PARAMETER(CondParam);

class CondState {

public:

CondParam params;

CachedOpPtr cond_op;

LoopState then_branch;

LoopState else_branch;

int branch_selection; // 1 if then branch; 0 if else branch; -1 if undefined

CondState(CondParam params,

const nnvm::Symbol& cond,

const nnvm::Symbol& then_sym,

const nnvm::Symbol& else_sym)

: params(std::move(params)),

cond_op(LoopState::MakeSharedOp(cond)),

then_branch(then_sym),

else_branch(else_sym),

branch_selection(-1) {}

};

static void CondComputeExCPU(const OpStatePtr& state_ptr,

const OpContext& ctx,

const std::vector<NDArray>& inputs,

const std::vector<OpReqType>& req,

const std::vector<NDArray>& outputs) {

// The argument `inputs' are loop_vars and other inputs

// loop_vars are stored in stored in `loop_vars_locs'

// The argument `outputs' are output and new_loop_vars

// [0: num_out_data) are outputs at each step.

// [num_out_data: ) are new_loop_vars

CondState& state = state_ptr.get_state<CondState>();

const CondParam& params = state.params;

// a helper function, converting std::vector<NDArray> to std::vector<NDArray*>

const auto to_ptr_vec = [](std::vector<NDArray>& in, std::vector<NDArray*>* out) {

out->clear();

out->reserve(in.size());

std::transform(

std::begin(in), std::end(in), std::back_inserter(*out), [](NDArray& a) { return &a; });

};

// sanity checks

CHECK_EQ(inputs.size() + 3U, (size_t)params.num_args);

CHECK_EQ(outputs.size(), (size_t)params.num_outputs);

CHECK_EQ(outputs.size(), req.size());

// construct inputs and outputs for cond

std::vector<NDArray> cond_inputs;

std::vector<NDArray> cond_outputs = {NDArray()};

std::vector<NDArray*> cond_input_ptr;

std::vector<NDArray*> cond_output_ptr;

extract_by_loc(inputs, params.cond_input_locs, &cond_inputs);

to_ptr_vec(cond_inputs, &cond_input_ptr);

to_ptr_vec(cond_outputs, &cond_output_ptr);

int& branch_selection = state.branch_selection;

// run cond

CHECK(cond_input_ptr.size() > 0) << "condition requires at least 1 input";

Context default_ctx = cond_inputs[0].ctx();

state.cond_op->Forward(nullptr, cond_input_ptr, cond_output_ptr, default_ctx);

branch_selection = as_bool_scalar(*cond_output_ptr[0]);

// select the right branch

const mxnet::Tuple<dim_t>& func_input_locs =

branch_selection ? params.then_input_locs : params.else_input_locs;

LoopState& loop_state = branch_selection ? state.then_branch : state.else_branch;

// extract inputs for the branch

std::vector<NDArray> func_inputs;

extract_by_loc(inputs, func_input_locs, &func_inputs);

loop_state.Forward(0, func_inputs, req, outputs, ctx.need_grad);

}

static void CondGradComputeExCPU(const OpStatePtr& state_ptr,

const OpContext& ctx,

const std::vector<NDArray>& inputs,

const std::vector<OpReqType>& _req,

const std::vector<NDArray>& outputs) {

CondState& state = state_ptr.get_state<CondState>();

const CondParam& params = state.params;

// sanity checks

CHECK_EQ(outputs.size() + 3U, (size_t)params.num_args);

CHECK_EQ(outputs.size(), _req.size());

// select the right branch

int branch_selection = state.branch_selection;

CHECK_NE(branch_selection, -1);

const mxnet::Tuple<dim_t>& func_input_locs =

branch_selection ? params.then_input_locs : params.else_input_locs;

LoopState& loop_state = branch_selection ? state.then_branch : state.else_branch;

// construct parameters

std::vector<NDArray> ograds(inputs.begin(), inputs.begin() + params.num_outputs);

std::vector<OpReqType> req;

extract_by_loc(_req, func_input_locs, &req);

std::vector<NDArray> igrads;

extract_by_loc(outputs, func_input_locs, &igrads);

loop_state.Backward(0, ograds, req, igrads);

loop_state.Cleanup();

}

static bool CondType(const nnvm::NodeAttrs& attrs,

std::vector<int>* in_type,

std::vector<int>* out_type) {

const CondParam& params = nnvm::get<CondParam>(attrs.parsed);

static const std::function<bool(const int&)> is_udf = is_type_udf;

CHECK_EQ(in_type->size() + 3U, (size_t)params.num_args);

CHECK_EQ(out_type->size(), (size_t)params.num_outputs);

CHECK_EQ(attrs.subgraphs.size(), 3U);

CHECK_EQ(attrs.subgraphs[0]->outputs.size(), 1U);

CHECK_EQ(attrs.subgraphs[1]->outputs.size(), attrs.subgraphs[2]->outputs.size());

std::vector<int> cond_in_type;

std::vector<int> then_in_type;

std::vector<int> else_in_type;

extract_by_loc(*in_type, params.cond_input_locs, &cond_in_type);

extract_by_loc(*in_type, params.then_input_locs, &then_in_type);

extract_by_loc(*in_type, params.else_input_locs, &else_in_type);

std::vector<int> cond_out_type = {0};

bool succ_0 = InferSubgraphDataType(*attrs.subgraphs[0], &cond_in_type, &cond_out_type);

CHECK(sync_in_in(params.cond_input_locs, in_type, &cond_in_type, is_udf));

bool succ_1 = InferSubgraphDataType(*attrs.subgraphs[1], &then_in_type, out_type);

CHECK(sync_in_in(params.then_input_locs, in_type, &then_in_type, is_udf));

bool succ_2 = InferSubgraphDataType(*attrs.subgraphs[2], &else_in_type, out_type);

CHECK(sync_in_in(params.else_input_locs, in_type, &else_in_type, is_udf));

return succ_0 && succ_1 && succ_2;

}

static bool CondStorageType(const nnvm::NodeAttrs& attrs,

const int dev_mask,

DispatchMode* dispatch_mode,

std::vector<int>* in_attrs,

std::vector<int>* out_attrs) {

const CondParam& params = nnvm::get<CondParam>(attrs.parsed);

static const std::function<bool(const int&)> is_udf = is_stype_udf;

CHECK_EQ(in_attrs->size() + 3U, (size_t)params.num_args);

CHECK_EQ(out_attrs->size(), (size_t)params.num_outputs);

CHECK_EQ(attrs.subgraphs.size(), 3U);

CHECK_EQ(attrs.subgraphs[0]->outputs.size(), 1U);

CHECK_EQ(attrs.subgraphs[1]->outputs.size(), attrs.subgraphs[2]->outputs.size());

std::vector<int> cond_in_attrs;

std::vector<int> then_in_attrs;

std::vector<int> else_in_attrs;

extract_by_loc(*in_attrs, params.cond_input_locs, &cond_in_attrs);

extract_by_loc(*in_attrs, params.then_input_locs, &then_in_attrs);

extract_by_loc(*in_attrs, params.else_input_locs, &else_in_attrs);

std::vector<int> cond_out_attrs = {kDefaultStorage};

DispatchMode cond_mode = DispatchMode::kUndefined;

DispatchMode then_mode = DispatchMode::kUndefined;

DispatchMode else_mode = DispatchMode::kUndefined;

*dispatch_mode = DispatchMode::kFComputeEx;

bool succ_0 = InferSubgraphStorage(

*attrs.subgraphs[0], dev_mask, &cond_mode, &cond_in_attrs, &cond_out_attrs);

CHECK(sync_in_in(params.cond_input_locs, in_attrs, &cond_in_attrs, is_udf));

bool succ_1 =