<jupyter_start><jupyter_text>Example: optimize a tensor of values so that their rank matches a target<jupyter_code>import torch
from perturbations_torch.fenchel_young import FenchelYoungLoss


def ranks(inputs, dim=-1):
    """Returns the ranks of the input values among the given axis."""
    return 1 + inputs.argsort(dim).argsort(dim).type(inputs.dtype)

# We initialize a random tensor
x = torch.randn([3, 5]).float()
print(x)

# Turn its grad on, since we will change this tensor to minimize our loss
x.requires_grad = True
y_true = torch.arange(5).float().unsqueeze(0).repeat([x.shape[0], 1])

print("Initially, the values in our tensor do not result in the desired argsort")
print(x.argsort(-1))
# Initialize an SGD optimizer and do 200 steps
optim = torch.optim.SGD([x], 0.01)

for iteration in range(200):
    optim.zero_grad()
    criterion = FenchelYoungLoss(ranks)
    loss = criterion(y_true, x).sum()
    loss.backward()
    optim.step()
    if iteration % 50 == 0:
        print(x.argsort(-1))
        print(loss.item())
        
print("SGD has succesfully changed our tensor to match the desired argsort!")<jupyter_output>tensor([[[1, 4, 2, 3, 0],
         [0, 4, 2, 1, 3],
         [4, 0, 1, 2, 3]]])
75.0
tensor([[[1, 0, 2, 3, 4],
         [0, 4, 1, 2, 3],
         [0, 4, 1, 2, 3]]])
41.00177764892578
tensor([[[0, 1, 2, 3, 4],
         [0, 1, 2, 3, 4],
         [0, 1, 2, 4, 3]]])
20.642358779907227
tensor([[[0, 1, 2, 3, 4],
         [0, 1, 2, 3, 4],
         [0, 1, 2, 4, 3]]])
17.00006675720215
SGD has succesfully changed our tensor to match the desired argsort!
{'fraction_non_alphanumeric': 0.15258511979823455, 'fraction_numerical': 0.08448928121059268, 'mean_word_length': 3.3008130081300813, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '521567', 'license': 'Apache-2.0', 'n_tokens_mistral': 674, 'n_tokens_neox': 572, 'n_words': 191, 'path': 'example-ranking-loss.ipynb', 'repo_name': 'nlml/perturbations-torch', 'size': '1586'}