Clean tests
Continue cleanup
Cleanup softmax