docs: private static IOURet calculateIOULabelPredicted(INDArray labelTL, INDArray labelBR, INDArray predictedWH, INDArray predictedXYinGridBox, INDArray objectPresentMask, INDArray objectPresentMaskBool){
          // FIXME: int cast
          int mb = (int) labelTL.size(0);
          int h = (int) labelTL.size(2);
          int w = (int) labelTL.size(3);
          int b = (int) predictedWH.size(1);
  
          INDArray labelWH = labelBR.sub(labelTL);                //4d [mb, 2, H, W], label W/H in terms of number of grid boxes
  
          int gridH = (int) labelTL.size(2);
          int gridW = (int) labelTL.size(3);
          //Add grid positions to the predicted XY values (to get predicted XY in terms of grid cell units in image,
          // from (0 to 1 in grid cell) format)
          INDArray linspaceX = Nd4j.linspace(0, gridW-1, gridW, predictedWH.dataType());
          INDArray linspaceY = Nd4j.linspace(0, gridH-1, gridH, predictedWH.dataType());
          INDArray grid = Nd4j.createUninitialized(predictedWH.dataType(), new long[]{2, gridH, gridW}, 'c');
          INDArray gridX = grid.get(point(0), all(), all());
          INDArray gridY = grid.get(point(1), all(), all());
          Broadcast.copy(gridX, linspaceX, gridX, 1);
          Broadcast.copy(gridY, linspaceY, gridY, 0);
  
          //Calculate X/Y position overall (in grid box units) from "position in current grid box" format
          INDArray predictedXY = predictedXYinGridBox.ulike();;
          Broadcast.add(predictedXYinGridBox, grid, predictedXY, 2,3,4); // [2, H, W] to [mb, b, 2, H, W]
  
  
          INDArray halfWH = predictedWH.mul(0.5);
          INDArray predictedTL_XY = halfWH.rsub(predictedXY);     //xy - 0.5 * wh
          INDArray predictedBR_XY = halfWH.add(predictedXY);      //xy + 0.5 * wh
  
          INDArray maxTL = predictedTL_XY.ulike();   //Shape: [mb, b, 2, H, W]
          Broadcast.max(predictedTL_XY, labelTL, maxTL, 0, 2, 3, 4);
          INDArray minBR = predictedBR_XY.ulike();
          Broadcast.min(predictedBR_XY, labelBR, minBR, 0, 2, 3, 4);
  
          INDArray diff = minBR.sub(maxTL);
          INDArray intersectionArea = diff.prod(2);   //[mb, b, 2, H, W] to [mb, b, H, W]
          Broadcast.mul(intersectionArea, objectPresentMask, intersectionArea, 0, 2, 3);
  
          //Need to mask the calculated intersection values, to avoid returning non-zero values when intersection is actually 0
          //No intersection if: xP + wP/2 < xL - wL/2 i.e., BR_xPred < TL_xLab   OR  TL_xPred > BR_xLab (similar for Y axis)
          //Here, 1 if intersection exists, 0 otherwise. This is doing x/w and y/h simultaneously
          INDArray noIntMask1 = Nd4j.createUninitialized(DataType.BOOL, maxTL.shape(), maxTL.ordering());
          INDArray noIntMask2 = Nd4j.createUninitialized(DataType.BOOL, maxTL.shape(), maxTL.ordering());
          //Does both x and y on different dims
          Broadcast.lte(predictedBR_XY, labelTL, noIntMask1, 0, 2, 3, 4);  //Predicted BR <= label TL
          Broadcast.gte(predictedTL_XY, labelBR, noIntMask2, 0, 2, 3, 4);  //predicted TL >= label BR
  
          noIntMask1 = Transforms.or(noIntMask1.get(all(), all(), point(0), all(), all()), noIntMask1.get(all(), all(), point(1), all(), all()) );    //Shape: [mb, b, H, W]. Values 1 if no intersection
          noIntMask2 = Transforms.or(noIntMask2.get(all(), all(), point(0), all(), all()), noIntMask2.get(all(), all(), point(1), all(), all()) );
          INDArray noIntMask = Transforms.or(noIntMask1, noIntMask2 );
  
          INDArray intMask = Transforms.not(noIntMask); //Values 0 if no intersection
          Broadcast.mul(intMask, objectPresentMaskBool, intMask, 0, 2, 3);
  
          //Mask the intersection area: should be 0 if no intersection
          intMask = intMask.castTo(predictedWH.dataType());
          intersectionArea.muli(intMask);
  
  
          //Next, union area is simple: U = A1 + A2 - intersection
          INDArray areaPredicted = predictedWH.prod(2);   //[mb, b, 2, H, W] to [mb, b, H, W]
          Broadcast.mul(areaPredicted, objectPresentMask, areaPredicted, 0,2,3);
          INDArray areaLabel = labelWH.prod(1);           //[mb, 2, H, W] to [mb, H, W]
  
          INDArray unionArea = Broadcast.add(areaPredicted, areaLabel, areaPredicted.dup(), 0, 2, 3);
          unionArea.subi(intersectionArea);
          unionArea.muli(intMask);
  
          INDArray iou = intersectionArea.div(unionArea);
          BooleanIndexing.replaceWhere(iou, 0.0, Conditions.isNan()); //0/0 -> NaN -> 0
  
          //Apply the "object present" mask (of shape [mb, h, w]) - this ensures IOU is 0 if no object is present
          Broadcast.mul(iou, objectPresentMask, iou, 0, 2, 3);
  
          //Finally, calculate derivatives:
          INDArray maskMaxTL = Nd4j.createUninitialized(DataType.BOOL, maxTL.shape(), maxTL.ordering());    //1 if predicted Top/Left is max, 0 otherwise
          Broadcast.gt(predictedTL_XY, labelTL, maskMaxTL, 0, 2, 3, 4);   // z = x > y
          maskMaxTL = maskMaxTL.castTo(predictedWH.dataType());
  
          INDArray maskMinBR = Nd4j.createUninitialized(DataType.BOOL, maxTL.shape(), maxTL.ordering());    //1 if predicted Top/Left is max, 0 otherwise
          Broadcast.lt(predictedBR_XY, labelBR, maskMinBR, 0, 2, 3, 4);   // z = x < y
          maskMinBR = maskMinBR.castTo(predictedWH.dataType());
  
          //dI/dx = lambda * (1^(min(x1+w1/2) - 1^(max(x1-w1/2))
          //dI/dy = omega * (1^(min(y1+h1/2) - 1^(max(y1-h1/2))
          //omega = min(x1+w1/2,x2+w2/2) - max(x1-w1/2,x2+w2/2)       i.e., from diff = minBR.sub(maxTL), which has shape [mb, b, 2, h, w]
          //lambda = min(y1+h1/2,y2+h2/2) - max(y1-h1/2,y2+h2/2)
          INDArray dI_dxy = maskMinBR.sub(maskMaxTL);              //Shape: [mb, b, 2, h, w]
          INDArray dI_dwh = maskMinBR.addi(maskMaxTL).muli(0.5);    //Shape: [mb, b, 2, h, w]
  
          dI_dxy.get(all(), all(), point(0), all(), all()).muli(diff.get(all(), all(), point(1), all(), all()));
          dI_dxy.get(all(), all(), point(1), all(), all()).muli(diff.get(all(), all(), point(0), all(), all()));
  
          dI_dwh.get(all(), all(), point(0), all(), all()).muli(diff.get(all(), all(), point(1), all(), all()));
          dI_dwh.get(all(), all(), point(1), all(), all()).muli(diff.get(all(), all(), point(0), all(), all()));
  
          //And derivatives WRT IOU:
          INDArray uPlusI = unionArea.add(intersectionArea);
          INDArray u2 = unionArea.mul(unionArea);
          INDArray uPlusIDivU2 = uPlusI.div(u2);   //Shape: [mb, b, h, w]
          BooleanIndexing.replaceWhere(uPlusIDivU2, 0.0, Conditions.isNan());     //Handle 0/0
  
          INDArray dIOU_dxy = Nd4j.createUninitialized(predictedWH.dataType(), new long[]{mb, b, 2, h, w}, 'c');
          Broadcast.mul(dI_dxy, uPlusIDivU2, dIOU_dxy, 0, 1, 3, 4);   //[mb, b, h, w] x [mb, b, 2, h, w]
  
          INDArray predictedHW = Nd4j.createUninitialized(predictedWH.dataType(), new long[]{mb, b, 2, h, w}, predictedWH.ordering());
          //Next 2 lines: permuting the order... WH to HW along dimension 2
          predictedHW.get(all(), all(), point(0), all(), all()).assign(predictedWH.get(all(), all(), point(1), all(), all()));
          predictedHW.get(all(), all(), point(1), all(), all()).assign(predictedWH.get(all(), all(), point(0), all(), all()));
  
          INDArray Ihw = predictedHW.ulike();;
          Broadcast.mul(predictedHW, intersectionArea, Ihw, 0, 1, 3, 4 );    //Predicted_wh: [mb, b, 2, h, w]; intersection: [mb, b, h, w]
  
          INDArray dIOU_dwh = Nd4j.createUninitialized(predictedHW.dataType(), new long[]{mb, b, 2, h, w}, 'c');
          Broadcast.mul(dI_dwh, uPlusI, dIOU_dwh, 0, 1, 3, 4);
          dIOU_dwh.subi(Ihw);
          Broadcast.div(dIOU_dwh, u2, dIOU_dwh, 0, 1, 3, 4);
          BooleanIndexing.replaceWhere(dIOU_dwh, 0.0, Conditions.isNan());     //Handle division by 0 (due to masking, etc)
  
          return new IOURet(iou, dIOU_dxy, dIOU_dwh);
      }