Switch the softmax to use the new deterministic reductions on the GPU,
results in a speed up of 10-40x on the existing ImageNet benchmarks and 2-3x on the newly added transformer benchmarks. Update the benchmark to also run on the GPU. Remove duplicate cpu tests. PiperOrigin-RevId: 168596693
Loading
Please sign in to comment