NarrowEmbedding.Config(
        channels_in=2,
        channels_out=4,
        dtype=torch.bfloat16,
        inner=Embedding.Config(
            channels_in=2,
            channels_out=4,
            padding_idx=None,
            device=None,
            dtype=None,
            lookup_dtype=None,
            shard=None,
            depth_index=(),
            init_weight=functools.partial(priml.model.init.truncated_normal, std=0.02)
        )
    )
