code: def test_pytorch_yolo_fragment():
      if not is_backend_tested("torch"):
          pytest.skip()
  
      import torch
  
      def old_way(
          input, num_classes, num_anchors, anchors, stride_h, stride_w
      ):
          # https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L28-L44
          bs = input.size(0)
          in_h = input.size(2)
          in_w = input.size(3)
          scaled_anchors = [
              (a_w / stride_w, a_h / stride_h) for a_w, a_h in anchors
          ]
  
          prediction = (
              input.view(bs, num_anchors, 5 + num_classes, in_h, in_w)
              .permute(0, 1, 3, 4, 2)
              .contiguous()
          )
          # Get outputs
          x = torch.sigmoid(prediction[..., 0])  # Center x
          y = torch.sigmoid(prediction[..., 1])  # Center y
          w = prediction[..., 2]  # Width
          h = prediction[..., 3]  # Height
          conf = torch.sigmoid(prediction[..., 4])  # Conf
          pred_cls = torch.sigmoid(prediction[..., 5:])  # Cls pred.
  
          # https://github.com/BobLiu20/YOLOv3_PyTorch/blob/c6b483743598b5f64d520d81e7e5f47ba936d4c9/nets/yolo_loss.py#L70-L92
          FloatTensor = (
              torch.cuda.FloatTensor if x.is_cuda else torch.FloatTensor
          )
          LongTensor = (
              torch.cuda.LongTensor if x.is_cuda else torch.LongTensor
          )
          # Calculate offsets for each grid
          grid_x = (
              torch.linspace(0, in_w - 1, in_w)
              .repeat(in_w, 1)
              .repeat(bs * num_anchors, 1, 1)
              .view(x.shape)
              .type(FloatTensor)
          )
          grid_y = (
              torch.linspace(0, in_h - 1, in_h)
              .repeat(in_h, 1)
              .t()
              .repeat(bs * num_anchors, 1, 1)
              .view(y.shape)
              .type(FloatTensor)
          )
          # Calculate anchor w, h
          anchor_w = FloatTensor(scaled_anchors).index_select(
              1, LongTensor([0])
          )
          anchor_h = FloatTensor(scaled_anchors).index_select(
              1, LongTensor([1])
          )
          anchor_w = (
              anchor_w.repeat(bs, 1)
              .repeat(1, 1, in_h * in_w)
              .view(w.shape)
          )
          anchor_h = (
              anchor_h.repeat(bs, 1)
              .repeat(1, 1, in_h * in_w)
              .view(h.shape)
          )
          # Add offset and scale with anchors
          pred_boxes = FloatTensor(prediction[..., :4].shape)
          pred_boxes[..., 0] = x.data + grid_x
          pred_boxes[..., 1] = y.data + grid_y
          pred_boxes[..., 2] = torch.exp(w.data) * anchor_w
          pred_boxes[..., 3] = torch.exp(h.data) * anchor_h
          # Results
          _scale = torch.Tensor([stride_w, stride_h] * 2).type(
              FloatTensor
          )
          output = torch.cat(
              (
                  pred_boxes.view(bs, -1, 4) * _scale,
                  conf.view(bs, -1, 1),
                  pred_cls.view(bs, -1, num_classes),
              ),
              -1,
          )
          return output
  
      def new_way(
          input, num_classes, num_anchors, anchors, stride_h, stride_w
      ):
          raw_predictions = rearrange(
              input,
              " b (anchor prediction) h w -> prediction b anchor h w",
              anchor=num_anchors,
          )
  
          anchors = torch.FloatTensor(anchors).to(input.device)
          anchor_sizes = rearrange(
              anchors, "anchor dim -> dim () anchor () ()"
          )
  
          _, _, _, in_h, in_w = raw_predictions.shape
          grid_h = rearrange(
              torch.arange(in_h).float(), "h -> () () h ()"
          ).to(input.device)
          grid_w = rearrange(
              torch.arange(in_w).float(), "w -> () () () w"
          ).to(input.device)
  
          predicted_bboxes = torch.zeros_like(raw_predictions)
          predicted_bboxes[0] = (
              raw_predictions[0].sigmoid() + grid_h
          ) * stride_h  # center y
          predicted_bboxes[1] = (
              raw_predictions[1].sigmoid() + grid_w
          ) * stride_w  # center x
          predicted_bboxes[2:4] = (
              raw_predictions[2:4].exp()
          ) * anchor_sizes  # bbox width and height
          predicted_bboxes[4] = raw_predictions[
              4
          ].sigmoid()  # confidence
          predicted_bboxes[5:] = raw_predictions[
              5:
          ].sigmoid()  # class predictions
          # only to match results of original code, not needed
          return rearrange(
              predicted_bboxes,
              "prediction b anchor h w -> b anchor h w prediction",
          )
  
      stride_h = 4
      stride_w = 4
      batch_size = 5
      num_classes = 12
      anchors = [[50, 100], [100, 50], [75, 75]]
      num_anchors = len(anchors)
  
      input = torch.randn(
          [batch_size, num_anchors * (5 + num_classes), 1, 1]
      )
      result1 = old_way(
          input=input,
          num_anchors=num_anchors,
          num_classes=num_classes,
          stride_h=stride_h,
          stride_w=stride_w,
          anchors=anchors,
      )
      result2 = new_way(
          input=input,
          num_anchors=num_anchors,
          num_classes=num_classes,
          stride_h=stride_h,
          stride_w=stride_w,
          anchors=anchors,
      )
      result1 = result1.reshape(result2.shape)
      assert torch.allclose(result1, result2)
  
file_path: einops/tests/test_examples.py
imports: import numpy
  import pytest
  from einops import rearrange, parse_shape, reduce
  from tests import is_backend_tested
  from tests.test_ops import imp_op_backends
label_desc: Use image node_id 3 for calling a global function with example usage: test_pytorch_yolo_fragment() and returns: output, rearrange