[Relax] Allow softmax to work on a large tensor when dimension is not the last one

[hugolatendresse]

For large tensors with non-last dimension softmax, we transpose to move the softmax dimension to the end, apply softmax, and then transpose back to the original shape.

Since the issue here that using softmax on non-last dimension could cause python/tvm/dlight/gpu/general_reduction.py to create arrays that are too big for the GPU shared memory, I tried to address this TODO by making changes to general_reduction.py, without success. However, as I was experimenting, I added a suggested handling for the case where num_leading_s = 0 in general_reduction.py. I thought I might as well leave that in the PR.

cc: @MasterJH5574

Edit: we may not merge this at all because it's better to fix the reduction directly, and the fix in this PR may simply be extra overhead

[tqchen]

let us instead to work and allow dlight to work correctly for non-last dimension cases

[hugolatendresse]

let us instead to work and allow dlight to work correctly for non-last dimension cases

Sounds good, closing the PR

[MasterJH5574]

Fixed in #17754