def scaled_loss(x, fn, a):
    """Apply a loss function to a tensor and pre- and post-scale it.
    Args:
        x: the data tensor, should already be squared: `x = y**2`.
        fn: the loss function, with signature `fn(x) -> y`.
        a: the scale parameter.
    Returns:
        The value of the loss, and its first and second derivatives.
    """
    a2 = a**2
    loss, loss_d1, loss_d2 = fn(x/a2)
    return loss*a2, loss_d1, loss_d2/a2