zheng-da commented on a change in pull request #10451: [MXNET-432] Add Foreach
URL: https://github.com/apache/incubator-mxnet/pull/10451#discussion_r189734486
##########
File path: python/mxnet/ndarray/contrib.py
##########
@@ -96,3 +98,101 @@ def rand_zipfian(true_classes, num_sampled, range_max,
ctx=None):
expected_count_sampled = expected_prob_sampled * num_sampled
return sampled_classes, expected_count_true, expected_count_sampled
# pylint: enable=line-too-long
+
+def foreach(func, data, init_states):
+ """Run a for loop with user-defined computation over NDArrays on dimension
0.
+
+ This operator simulates a for loop and func has the computation for an
iteration
+ of the for loop. It runs the computation in func on each slice from the
input
+ NDArrays.
+
+ func takes two arguments as input and outputs a tuple of two elements,
+ as illustrated below:
+
+ out, states = func(data1, states)
+
+ data1 can be either an NDArray or a list of NDArrays. If data is an
NDArray,
+ data1 is an NDArray. Otherwise, data1 is a list of NDArrays and has the
same
+ size as data. states is a list of NDArrays and have the same size as
init_states.
+ Similarly, out can be either an NDArray or a list of NDArrays, which are
concatenated
+ as the first output of foreach; states from the last execution of func
+ are the second output of foreach.
+
+ The computation done by this operator is equivalent to the pseudo code
below
+ when the input data is NDArray:
+
+ states = init_states
+ outs = []
+ for i in data.shape[0]:
+ s = data[i]
+ out, states = func(s, states)
+ outs.append(out)
+ outs = stack(*outs)
+
+
+ Parameters
+ ----------
+ func : a Python function.
Review comment:
this is to follow the interface of TensorFlow.
https://www.tensorflow.org/api_docs/python/tf/while_loop
Using class does make API more well defined, but it requires users to write
more code to define it. I don't know what is the best way.
@piiswrong what's your opinion?
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
