def _get_layer_lrs(learning_rate, layer_decay, n_layers=2):
    """Have lower learning rates for layers closer to the input."""
    key_to_depths = {}
    
    key_to_depths["layer_0/"] = 0
    key_to_depths["layer_1/"] = 4

    print(key_to_depths)

    return {
        key: learning_rate * (layer_decay ** (n_layers + 2 - depth))
        for key, depth in key_to_depths.items()
    }