Method init

axlearn/common/repeat.py:154–161

(prng_key_i)

        self, prng_key: Tensor, *, prebuilt: Optional[Nested[Optional[ParameterSpec]]] = None
    ) -> NestedTensor:
        def init(prng_key_i):
            return VDict(
                layer=self.layer.initialize_parameters_recursively(
                    # `prebuilt` must be consistent across all layers.
                    prng_key_i,
                    prebuilt=get_or_none(prebuilt, "layer"),
                )
            )

        cfg = self.config
        return jax.vmap(init)(split_prng_key(prng_key, cfg.num_layers).keys)