code: public static INDArray unitVec(INDArray toScale) {
          double length = toScale.norm2Number().doubleValue();
  
          if (length > 0) {
              if (toScale.data().dataType() == (DataType.FLOAT))
                  return Nd4j.getBlasWrapper().scal(1.0f / (float) length, toScale);
              else
                  return Nd4j.getBlasWrapper().scal(1.0 / length, toScale);
  
          }
          return toScale;
      }
docs: Scale by 1 / norm2 of the matrix
  
  @param toScale the ndarray to scale
  @return the scaled ndarray