def t5x_mlp_lookup(params, i, prefix, split_mlp_wi=False):
    """Returns the MLP parameters of a layer. Does not transpose."""
    if split_mlp_wi:
        wi_0 = params[f"{prefix}/{prefix}/mlp/wi_0/kernel"][:, i, :]
        wi_1 = params[f"{prefix}/{prefix}/mlp/wi_1/kernel"][:, i, :]
        wi = (wi_0, wi_1)
    else:
        wi = params[f"{prefix}/{prefix}/mlp/wi/kernel"][:, i, :]

    wo = params[f"{prefix}/{prefix}/mlp/wo/kernel"][:, i, :]
    return wi, wo
Returns the MLP parameters of a layer. Does not transpose.
dpkd/transformers/src/transformers/models/umt5/convert_umt5_checkpoint_to_pytorch.py