How to set an Module attribute after constructing it?

[marcvanzee]

Original question by @salayatana66:

I want a unit_test that checks whether a certain submodule gets called. I have something like this:

bound_transformer = Transformer(self.trans_cfg, self.aspect_vocab_size,
                                    AdditiveEmbeddingsApplicationMode.ADD, self.start_layer_tail,
                                    parent=flax.core.Scope(params))
decoder_head = bound_transformer.decoder_head
decoder_head.decode_embedded_targets = mock.Mock(side_effect=decoder_head.decode_embedded_targets)
...

but then Flax complains at decoder_head.decode_embed_targets = ... by Module instance is frozen outside of setup method; is there a way I can get around this?

[marcvanzee]

You can use object.__setattr__:

object.__setattr__(decoder_head, 'decode_embedded_targets', mock.Mock(side_effect=decoder_head.decode_embedded_targets)