How to use apply with additional parameters?

[Lando-L]

I am having trouble understanding the apply method when used with additional parameters, such as with the hk.nets.MLP module.

I tried to implement a simple CNN architecture of several blocks of convolutions followed by a few layers of fully connected layers. Describing the network architecture and initialising the parameters works as expected.

class SimpleCNN(hk.Module):
    def __init__(self):
        self.cnn_blocks = ...
        self.fc_layers = hk.MLP([...])

    def __call__(self, inputs, dropout_rate, rng):
        return self.fc_layers(self.cnn_blocks(inputs), dropout_rate, rng)

def forward(inputs, dropout_rate, rng):
    return SimpleCNN()(inputs, dropout_rate, rng)

net = hk.transform(forward)
params = net.init(sample_rng, sample, None, None)

When it comes to applying the network I am getting confused. As I understand it, the apply methods takes the parameters params and rng, followed by the parameters the untransformed function is expecting. In the case of SimpleCNN that would be inputs, dropout_rate, and rng. Do I have to pass in two rng? If so, then what is the purpose of the first rng?

[Lando-L]

It looks like I missed haiku's random API.