Help understanding a bug when FoldList using RandomChoice is Compiled

Question

Bug introduced in 3.0, persisting through 13.1.

---

Preamble: I submitted this bug to [email protected] and it was acknowledged. I assume it will be fixed eventually. In the meantime I'd like to understand what's going on so I can have a better idea as to how concerned I should be about other code I've written that might be subject to the same or a similar issue.

I present a toy example that illustrates the bug. In the toy example I don't really need the compiled version because the uncompiled version runs faster, but my actual applications are much larger and more complicated and the compiled version runs faster.

Here is the uncompiled version that works correctly:

fun = Function[{},
  FoldList[
    RandomChoice[{{1, 0}, {0, 1}}[[#1]]*#2 -> {1, 2}] &,
    RandomChoice[{1, 1} -> {1, 2}],
    {{1, 1}, {1, 1}}
  ]
];

fun[] returns either {1,1,1} or {2,2,2} with equal probability as the following usage indicates:

KeySort@Counts@Table[fun[], 100]

(* <|{1, 1, 1} -> 52, {2, 2, 2} -> 48|> *)

The problem comes when the function is compiled, either explicitly or implicitly. To illustrate the effect of implicit compilation consider the following:

KeySort@Counts@Table[fun[], 250]

(* <|{1, 1, 1} -> 62, {1, 2, 2} -> 51, {2, 1, 1} -> 70, {2, 2, 2} -> 67|> *)

We see that erroneous output is produced. We can replicate this behavior by explicitly compiling the function:

cf = Compile[{},
  FoldList[
    RandomChoice[{{1, 0}, {0, 1}}[[#1]]*#2 -> {1, 2}] &,
    RandomChoice[{1, 1} -> {1, 2}],
    {{1, 1}, {1, 1}}
  ]
];

SeedRandom[1];
cf[]

(* {1, 2, 2}  *)

I have two questions: (1) What is going on? and (2) Why?

I don't know how to read the byte code in compiled functions, but the following indicates the presence of an extra call to RandomChoice:

Position[cf, "RandomChoiceWeights"]

(* {{6, 2, 2}, {6, 12, 2}, {6, 17, 2}} *)

So how the heck does Compile go about introducing this extra call? And how should I think about other uses of Compile that involve randomness? I just don't have a sense for how bad this might be. I'd like the answer to be "Oh this isn't bad at all because it's triggered by this very specific thing that doesn't apply to anything else."

Any thoughts would be appreciated.

Update

FunctionCompile does not suffer from the bug.

fc = FunctionCompile[
  Function[{},
    FoldList[
      RandomChoice[{{1, 0}, {0, 1}}[[#1]]*#2 -> {1, 2}] &, 
      RandomChoice[{1, 1} -> {1, 2}], 
      {{1, 1}, {1, 1}}
    ]
  ]
]

KeySort@Counts@Table[fc[], 250]

(* <|{1, 1, 1} -> 129, {2, 2, 2} -> 121|> *)

Answer 1

Here' s a simpler example that demonstrates the same issue:

fun = Function[{}, FoldList[#1 &, RandomChoice[{1, 2}], {0, 0, 0}]];

KeySort@Counts@Table[fun[], 100]
(* <|{1, 1, 1, 1} -> 42, {2, 2, 2, 2} -> 58|> *)

KeySort@Counts@Table[fun[], 250]
(* <|{1, 1, 1, 1} -> 62, {1, 2, 2, 2} -> 72, {2, 1, 1, 1} -> 58, {2, 2, 2, 2} -> 58|> *)

cf = Compile @@ fun;

KeySort@Counts@Table[cf[], 100]
(* <|{1, 1, 1, 1} -> 25, {1, 2, 2, 2} -> 27, {2, 1, 1, 1} -> 24, {2, 2, 2, 2} -> 24|> *)

It appears that the first element is being randomly chosen independently of the other, repeated, value. The CompilePrint output gives a hint as to why:

Needs["CompiledFunctionTools`"]

CompilePrint[cf]

< some stuff omitted >
1   I1 = RandomChoice[ T(I1)0]]
2   I6 = Length[ T(I1)1]
3   I4 = I7
4   T(I1)2 = Table[ I6]
5   I8 = I7
6   goto 8
7   Element[ T(I1)2, I4] = I1
8   if[ ++ I8 <= I6] goto 7
9   I2 = RandomChoice[ T(I1)0]]
10  T(I1)3 = {I5}
11  T(I2)4 = {T(I1)3}
12  T(I1)6 = Insert[ T(I1)2, T(I0)2, T(I2)4]]
13  Return

Recall that

FoldList[f,x,{a,b,...}] gives {x, f[x,a], f[f[x,a],b], ...}

The compiled function runs a loop (lines 7 & 8) which computes {f[x,a], f[f[x,a],b], ...} and then afterwards inserts x at the beginning of the list (line 12). The problem is that the RandomChoice is evaluated twice; once (line 1) for the loop and again (line 9) for the first element.

Obviously what should happen is for RandomChoice to evaluate just once, and the result used for both the loop and the first element.

As a workaround you can pull the RandomChoice out of the FoldList expression:

cf = Compile[{}, With[{r = RandomChoice[{1, 2}]},
    FoldList[#1 &, r, {0, 0, 0}]]];

KeySort@Counts@Table[cf[], 100]
(* <|{1, 1, 1, 1} -> 37, {2, 2, 2, 2} -> 63|> *)

It's hard to say how bad this is, I suspect it's hard to find other examples. Nevertheless it's worrying that the compiler doesn't fully respect the rules for evaluation/holding of function arguments, especially given that the compiler can be invisibly used "behind the scenes" e.g. in Table.

Answer 2

It seems FunctionCompile does not share the same issue and its result performs as good as that from Compile.

$Version
(* 13.1.0 for Microsoft Windows (64-bit) (June 16, 2022) *)

funcFC = FunctionCompile[
        FunctionDeclaration[foldcore,
            Typed[{"MachineInteger", "MachineInteger"} -> "MachineInteger"]@
                Function[{a, b}, a]
         ]
        , Function[{}, FoldList[foldcore, RandomChoice[{1, 2}], {0, 0, 0}]]
    ]; // AbsoluteTiming
(* Out[]= {1.55718, Null} *)

Table[funcFC[], 10^5] // Tally // AbsoluteTiming
(* Out[]= {0.0767217, {{{1, 1, 1, 1}, 50190}, {{2, 2, 2, 2}, 49810}}} *)

funcFC[]; // RepeatedTiming
(* Out[]= {7.52546*10^-7, Null} *)

As comparison, this is test result based on Compile with Simon's With trick:

funcC = Compile[{}, With[{r = RandomChoice[{1, 2}]}, FoldList[# &, r, {0, 0, 0}]]];

Table[funcC[], 10^5] // Tally // AbsoluteTiming
(* Out[]= {0.112121, {{{2, 2, 2, 2}, 50242}, {{1, 1, 1, 1}, 49758}}} *)

funcC[]; // RepeatedTiming
(* Out= {1.55156*10^-6, Null} *)

We can see from the IR that FunctionCompile does essentially what Simon's With trick does (the initial random choice is executed in line 5 and stored in %7, which is then reused in FoldList in line 6):

Needs["Compile`"]

pm = CompileToIR[
        FunctionDeclaration[foldcore,
            Typed[{"MachineInteger", "MachineInteger"} -> "MachineInteger"]@
                Function[{a, b}, a]
         ]
        , Function[{}, FoldList[foldcore, RandomChoice[{1, 2}], {0, 0, 0}]]
    ]

pm["getFunctionModule", "Main"]["toString"]

---

The original function in OP can be compiled similarly:

funcOP = FunctionCompile[
      FunctionDeclaration[foldcore,
        Typed[{"MachineInteger", "PackedArray"["MachineInteger", 1]} -> "MachineInteger"]@
            Function[{a, b},
                RandomChoice[{{1, 0}, {0, 1}}[[a]]*b -> {1, 2}]
             ]
       ]
      , Function[{}, FoldList[foldcore, RandomChoice[{1, 1} -> {1, 2}], {{1, 1}, {1, 1}}]]
  ]

Table[funcOP[], 10^5] // Tally // AbsoluteTiming
(* Out[]= {0.275108, {{{2, 2, 2}, 50044}, {{1, 1, 1}, 49956}}} *)

funcOP[]; // RepeatedTiming
(* Out[]= {3.03958*10^-6, Null} *)