# Can NetNestOperator take a non pre-defined number of iterations?

I'd like to Nest a particular NetChain inside a larger NetGraph for a random number of iterations per round or batch. Is this possible? Here are my (so far fruitless) attempts:

First, we define a dummy NetChain and illustrate how to Nest it for a fixed number of iterations using a few methods:

nc = NetInitialize[NetChain[{2, Ramp}, "Input" -> 2]];
nested[1] = NetNestOperator[nc, 3];
nested[2] = FunctionLayer[Nest[nc, #, 3] &];
nested[3] =
NetTake[NetGraph[FunctionLayer[NestList[nc, #, 3] &]],
3*Information[nc, "LayersCount"]];
nested[4] =
CompiledLayer[
FunctionCompile[
Function[Typed[arg, TypeSpecifier["PackedArray"]["Real64", 1]],
Nest[Typed[
KernelFunction[
nc], {TypeSpecifier["PackedArray"]["Real64", 1]} ->
TypeSpecifier["PackedArray"]["Real64", 1]], arg, 3]]]];

(*Show Equivalence*)
Through[(nested /@ Range[4])[{0.3, 0.2}]]
SameQ @@ %

Out[1]= {{0.369132, 0.}, {0.369132, 0.}, {0.369132, 0.}, {0.369132, 0.}}
Out[2]= True


From these, the last would require the additional definition of gradientfunc to be usable in NetTrain, so likely less appealing, and the third perhaps gives a hint on how this can be achieved semi-manually.

Naively, I hoped this syntax would allow for a second input port for the second method:

FunctionLayer[Function[{in, iter}, Nest[nc, in, iter]]]

During evaluation of In[3]:= FunctionLayer::compilerr: Cannot interpret Nest[NetChain[<2>], #1, #2] & as a network.
Out[3]= \$Failed


If this (or similar) worked, then one could attach a RandomArrayLayer as the second input, e.g. something like this:

ra=RandomArrayLayer[DiscreteUniformDistribution[{4,6}],"Output"->"Integer"];
ra[]

Out[4]= 6


### EDIT 01:

I guess one could do the following, where we're appending the random integer as the last entry of the input array (since CompiledLayer expects a single input):

nested[5] =
CompiledLayer[
FunctionCompile[
Function[Typed[arg, TypeSpecifier["PackedArray"]["Real64", 1]],
Nest[Typed[
KernelFunction[
nc], {TypeSpecifier["PackedArray"]["Real64", 1]} ->
TypeSpecifier["PackedArray"]["Real64", 1]], Most[arg],
Round[Last[arg]]]]]];

nested[5][{0.3, 0.2, 4}]
nested[5][{0.3, 0.2, 5}]
nested[5][{0.3, 0.2, 6}]

Out[5]= {0.336348, 0.}
Out[6]= {0.306477, 0.}
Out[7]= {0.279258, 0.}


But then I'm not sure how to compute gradientfunc..

### EDIT 02:

The accepted answer from @xslittlegrass does the trick very nicely! FWIW, much of that coded can be simplified as follows:

randomNesting = NetChain[{
FunctionLayer[NestList[nc, #, 9] &],
FunctionLayer[RandomChoice]
}]


We can use a particularly simple nc (a simple rotation matrix by 45°) to demonstrate this

nc = NetChain[{LinearLayer[2, "Weights" -> RotationMatrix[π/4.],
"Biases" -> None]}, "Input" -> 2];
Table[randomNesting[{0.3, 0.2}], 4]


{{-0.2, 0.3}, {-0.3, -0.2}, {-0.3, -0.2}, {-0.3, -0.2}}

## 1 Answer

I think one way to achieve this is to compute the all the nested results and then pick the correct one according the to random number, something like this

maxIteration = 10;
NestList[f, a, maxIter][[RandomInteger[{1, maxIter}]]]


This can be done by sharing the weights across all the repeated NetChain.

layers = Flatten@{
Table["nc" <> ToString[i] -> nc, {i, 1, maxIter}],
"cat" -> CatenateLayer[],
"reshape" -> ReshapeLayer[{maxIter, 2}],
"rand" ->
RandomArrayLayer[DiscreteUniformDistribution[{1, maxIter}],
"Output" -> {1}],
"part" -> ExtractLayer[]
};

connections = Flatten@{
Rule @@@
Partition[Table["nc" <> ToString[i], {i, 1, maxIter}], 2, 1],
Table["nc" <> ToString[i] -> "cat", {i, 1, maxIter}],
"rand" -> NetPort["part", "Position"],
"cat" -> "reshape" -> "part"
};

net = NetInitialize@NetGraph[layers, connections]


• This is great, thanks! FWIW, all of these can be achieved simply using FunctionLayer as follows: NetChain[{ FunctionLayer[NestList[nc, #, 9] &], FunctionLayer[RandomChoice] }]. I'll edit my question demonstrating this so it's easier for others to find it in the future! Commented Jul 12, 2021 at 1:17
• @GeorgeVarnavides That's nice! Good to know. Commented Jul 12, 2021 at 3:31