using scan with multiple layer - each with diffrent initialization - working example

[jakubMitura14]

The post is also published in flax forum https://github.com/google/flax/discussions/2912 - so this one can be removed by administrator

I used google/flax#2127 example a a basis
I want to have multiple layers of the same type but with diffrent value of a dummy variable.
Then sequentially invoke each layer, summing theit outputs to a single number.

Main problem is here is that when I pass array to LayerScanned I got the same array broadcasted to each layer instead of hetting diffrent inteer for each.
I also for some reason get the array as a final result not cumulative sum (although this is a minor problem)

import flax.linen as nn
import jax
import jax.numpy as jnp

class Layer(nn.Module):
  dummy :int
  @nn.compact
  def __call__(self, c, x, t):
    print(f"dummyy {self.dummy}")
    x = nn.Dense(len(x))(x)
    x = jax.nn.softmax(jnp.exp(t) * x)

    c = c + jnp.sum(jnp.ravel(x))
    return c, x

class Model(nn.Module):
  @nn.compact
  def __call__(self, x, t):
    LayerScanned = nn.scan(Layer,
                               variable_axes={
                                            "params": 0
                                            ,"dummy":0
                                            },
                           split_rngs={'params': False},
                           length=5,
                           in_axes=(0, nn.broadcast),
                           out_axes=1)
    carry = jnp.zeros_like(x)
    carry, x = LayerScanned(jnp.arange(5))(carry, x, t)
    return x,carry


x = jnp.zeros([5, 2])
model = Model()
params = model.init(jax.random.PRNGKey(0), x, 1)
print(f"params {params}")
res,c = model.apply(params, x,0.1)

[cgarciae]

Hey @jakubMitura14, I added some jax.debug.print statements to your code and can verify that:

• You do get different inputs x to each layer
• The carry c is accumulating the sums properly
• t is correctly being broadcasted

import flax.linen as nn
import jax
import jax.numpy as jnp


class Layer(nn.Module):
    dummy: int

    @nn.compact
    def __call__(self, c, x, t):
        jax.debug.print("x: {x}", x=x)
        jax.debug.print("t: {t}", t=t)
        x = nn.Dense(len(x))(x)
        x = jax.nn.softmax(jnp.exp(t) * x)

        c = c + jnp.sum(jnp.ravel(x))
        jax.debug.print("c: {c}\n", c=c.flatten())
        return c, x


class Model(nn.Module):
    @nn.compact
    def __call__(self, x, t):
        LayerScanned = nn.scan(
            Layer,
            variable_axes={"params": 0, "dummy": 0},
            split_rngs={"params": False},
            length=5,
            in_axes=(0, nn.broadcast),
            out_axes=0,
        )
        carry = jnp.zeros_like(x)
        carry, y = LayerScanned(10)(carry, x, t)
        jax.debug.print("final y: {y}", y=y)
        jax.debug.print("final c: {c}", c=carry)
        return y, carry


x = jax.random.uniform(jax.random.PRNGKey(0), (5, 2))
model = Model()
print("\ninit")
params = model.init(jax.random.PRNGKey(0), x, t=1.0)
print("\napply")
y, c = model.apply(params, x, t=0.1)

init
t: 1.0
x: [0.35490513 0.60419905]
t: 1.0
c: [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]

x: [0.4275843  0.23061597]
t: 1.0
c: [2. 2. 2. 2. 2. 2. 2. 2. 2. 2.]

x: [0.32985854 0.43953657]
t: 1.0
c: [3. 3. 3. 3. 3. 3. 3. 3. 3. 3.]

x: [0.25099766 0.27730572]
t: 1.0
c: [4. 4. 4. 4. 4. 4. 4. 4. 4. 4.]

x: [0.7678207  0.71474564]
t: 1.0
c: [5. 5. 5. 5. 5. 5. 5. 5. 5. 5.]

final y: [[0.15831344 0.8416866 ]
 [0.10467978 0.89532024]
 [0.17000116 0.8299988 ]
 [0.22761038 0.7723896 ]
 [0.0224549  0.97754514]]
final c: [[5. 5.]
 [5. 5.]
 [5. 5.]
 [5. 5.]
 [5. 5.]]

apply
t: 0.1
x: [0.35490513 0.60419905]
t: 0.10000000149011612
c: [1. 1. 1. 1. 1. 1. 1. 1. 1. 1.]

x: [0.4275843  0.23061597]
t: 0.10000000149011612
c: [2. 2. 2. 2. 2. 2. 2. 2. 2. 2.]

x: [0.32985854 0.43953657]
t: 0.10000000149011612
c: [3. 3. 3. 3. 3. 3. 3. 3. 3. 3.]

x: [0.25099766 0.27730572]
t: 0.10000000149011612
c: [4. 4. 4. 4. 4. 4. 4. 4. 4. 4.]

x: [0.7678207  0.71474564]
t: 0.10000000149011612
c: [5. 5. 5. 5. 5. 5. 5. 5. 5. 5.]

final y: [[0.33641553 0.66358453]
 [0.29471138 0.7052886 ]
 [0.34419286 0.65580714]
 [0.37830034 0.62169963]
 [0.17737925 0.8226208 ]]
final c: [[5. 5.]
 [5. 5.]
 [5. 5.]
 [5. 5.]
 [5. 5.]]

[jakubMitura14]

Thanks! However main problem is the dummy variable ( it is how is defined dummy: int) that is passed to the layer constructor during initialization!