docs: @Override
      public void onNDArrayPartial(INDArray arr, long idx, int... dimensions) {
          INDArray get = this.arr.get();
          get.tensorAlongDimension((int) idx, dimensions).assign(arr);
  
      }