NeuralLossFunction

class lightning_ir.loss.neural.NeuralLossFunction(temperature: float = 1, tol: float = 1e-05, max_iter: int = 50)[source]

Base class for neural loss functions that compute ranks from scores using neural sorting.

__init__(temperature: float = 1, tol: float = 1e-05, max_iter: int = 50) → None[source]

Initialize the NeuralLossFunction.

Parameters:

temperature (float) – Temperature parameter for scaling the scores. Defaults to 1.
tol (float) – Tolerance for convergence. Defaults to 1e-5.
max_iter (int) – Maximum number of iterations for convergence. Defaults to 50.

Methods

`__init__`([temperature, tol, max_iter])	Initialize the NeuralLossFunction.
`get_sorted_targets`(scores, targets)	Get the sorted targets based on the neural sort permutation matrix.
`neural_sort`(scores)	Compute the neural sort permutation matrix from scores.
`sinkhorn_scaling`(mat)	Apply Sinkhorn scaling to the permutation matrix.

get_sorted_targets(scores: Tensor, targets: Tensor) → Tensor[source]

Get the sorted targets based on the neural sort permutation matrix.

Parameters:

Returns:

The sorted targets tensor.

Return type:

torch.Tensor

neural_sort(scores: Tensor) → Tensor[source]

Compute the neural sort permutation matrix from scores.

sinkhorn_scaling(mat: Tensor) → Tensor[source]

Apply Sinkhorn scaling to the permutation matrix.