前向只选一个，反向为什么还有梯度？

"""Fixed-score straight-through arithmetic. Python standard library only."""
import json
import math
from pathlib import Path


def softmax(scores, temperature):
    if not scores or not math.isfinite(temperature) or temperature <= 0:
        raise ValueError('positive finite temperature and nonempty scores required')
    if not all(math.isfinite(x) for x in scores):
        raise ValueError('scores must be finite')
    maximum = max(scores)
    weights = [math.exp((x - maximum) / temperature) for x in scores]
    total = sum(weights)
    return [x / total for x in weights]


def dot(a, b):
    if len(a) != len(b):
        raise ValueError('different lengths')
    return sum(x * y for x, y in zip(a, b))


def one_hot(scores):
    chosen = max(range(len(scores)), key=scores.__getitem__)
    return [float(i == chosen) for i in range(len(scores))]


def surrogate_gradient(scores, costs, temperature, squared=False):
    probabilities = softmax(scores, temperature)
    average = dot(probabilities, costs)
    # This is an intentionally prescribed backward rule, not d(argmax)/dz.
    multiplier = 2 * dot(one_hot(scores), costs) if squared else 1.0
    return [multiplier * p * (cost - average) / temperature
            for p, cost in zip(probabilities, costs)]


def finite_difference(function, values, step=1e-6):
    result = []
    for index in range(len(values)):
        plus, minus = list(values), list(values)
        plus[index] += step
        minus[index] -= step
        result.append((function(plus) - function(minus)) / (2 * step))
    return result


def main():
    # Fixed noise g=(0,0,0): this isolates the backward rule, not sampling.
    logits, noise, costs = [math.log(2), 0.0, 0.0], [0.0] * 3, [1.0, 3.0, 5.0]
    scores = [z + g for z, g in zip(logits, noise)]
    probabilities = softmax(scores, 1.0)
    hard = one_hot(scores)
    linear_grad = surrogate_gradient(scores, costs, 1.0)
    hard_fd = finite_difference(lambda z: dot(one_hot(z), costs), scores)
    soft_fd = finite_difference(lambda z: dot(softmax(z, 1.0), costs), scores)
    squared_st = surrogate_gradient(scores, costs, 1.0, squared=True)
    squared_soft_fd = finite_difference(lambda z: dot(softmax(z, 1.0), costs) ** 2, scores)
    assert probabilities == [0.5, 0.25, 0.25]
    assert hard == [1.0, 0.0, 0.0] and hard_fd == [0.0] * 3
    assert max(abs(a-b) for a,b in zip(linear_grad, soft_fd)) < 1e-8
    assert linear_grad == [-0.75, 0.125, 0.625]
    assert squared_st == [-1.5, 0.25, 1.25]
    assert max(abs(a-b) for a,b in zip(squared_st, squared_soft_fd)) > 1
    temperatures = []
    for temperature in [0.5, 1.0, 2.0]:
        p = softmax(scores, temperature)
        gradient = surrogate_gradient(scores, costs, temperature)
        assert abs(sum(p) - 1) < 1e-12 and abs(sum(gradient)) < 1e-12
        assert one_hot(p) == hard
        temperatures.append({'temperature':temperature, 'soft':p, 'gradient':gradient})
    # A numerical derivative at a decision boundary is not a useful reference.
    tie = [0.0, 0.0, -2.0]
    tie_fd = finite_difference(lambda z: dot(one_hot(z), costs), tie)
    assert abs(tie_fd[0]) > 100000
    rejected = []
    for value in [0, -1, float('nan')]:
        try:
            softmax(scores, value)
        except ValueError:
            rejected.append(str(value))
    assert len(rejected) == 3
    report = {'scope':'fixed-noise arithmetic; no sampling expectation or model training',
        'logits':logits,'fixed_noise':noise,'costs':costs,'soft':probabilities,'hard':hard,
        'hard_linear_loss':dot(hard,costs),'soft_linear_loss':dot(probabilities,costs),
        'surrogate_linear_gradient':linear_grad,'hard_finite_difference':hard_fd,
        'soft_finite_difference':soft_fd,'squared_st_gradient':squared_st,
        'squared_soft_finite_difference':squared_soft_fd,'temperatures':temperatures,
        'tie_finite_difference':tie_fd,'rejected_temperatures':rejected,'checks_passed':True}
    Path('report.json').write_text(json.dumps(report, indent=2)+'\n',encoding='utf-8')
    print('soft =', probabilities, 'hard =', hard)
    print('linear loss: hard=1.000000 soft=2.500000')
    print('ST linear gradient =', linear_grad)
    print('hard finite difference =', hard_fd)
    print('soft finite difference =', [round(x, 6) for x in soft_fd])
    print('squared loss ST gradient =', squared_st)
    print('entirely soft squared gradient =', [round(x, 6) for x in squared_soft_fd])
    for row in temperatures:
        print('tau={:.1f} soft={} gradient={}'.format(row['temperature'],
            [round(x, 6) for x in row['soft']], [round(x, 6) for x in row['gradient']]))
    print('decision-boundary finite difference =', [round(x, 6) for x in tie_fd])
    print('checks=PASS; report.json written')


if __name__ == '__main__':
    main()