docs: public static SDVariable softmax(SameDiff SD, SDVariable x, int dimension, int rank){
          int[] permutation = ArrayUtil.range(0, rank);
          permutation[0] = dimension;
          permutation[dimension] = 0;
  
          return SD.nn.softmax(x.permute(permutation)).permute(ArrayUtil.invertPermutation(permutation));
      }