Problems with loop (body). TypeError: body_fun output and input must have identical types

[rietmann-nv]

Hello, I'm working on a scientific code, and having some trouble with building a loop with a somewhat longer loop body. I'm using the following looping features

with loops.Scope() as s:
    s.varA = jnp.zeros(...)
    ...
    for i in s.range(varOutside.shape[0]):
        loopBody with dynamic updates and index_updates (on s.varABC) and read-only on captured variables

I have some mutable variables in the loop which are in the s context variable, and some read-only variables which are captured in the loop scope.

I get the following error:

TypeError: body_fun output and input must have identical types, got
[ShapedArray(int32[]), ShapedArray(complex64[1,420,421]), ShapedArray(complex64[1,420,421]), ShapedArray(float32[192,192]), ShapedArray(complex64[192,192]), ShapedArray(float32[]), ShapedArray(float32[323,192,192])]
and
[ShapedArray(int32[], weak_type=True), ShapedArray(complex64[1,420,421]), ShapedArray(complex64[1,420,421]), ShapedArray(float32[192,192]), ShapedArray(float32[192,192]), ShapedArray(float32[]), ShapedArray(float32[323,192,192])].

The only difference I see, is the int32[] variable. One has weak_type=True and one does not. I'm not even sure which variable this int32[] corresponds to, and I can't seem to figure out a way to remove this error.

Can you give me some more examples of what I should be looking at?

Thanks!

[rietmann-nv]

Finally figured it out: I now noticed that one of the scoped variables changes from float32[...] to complex64[...], which was the actual cause of the problems. Writing this here in case anyone has this trouble in the future. Cheers!