Bug Report: macOS 15 Beta - PyTorch gridsample Not Utilising Apple Neural Engine on MacBook Pro M2

In macOS 15 beta the gridsample function from PyTorch is not executing as expected on the Apple Neural Engine in MacBook Pro M2.

Please find below a Python code snippet that demonstrates the problem:

import coremltools as ct
import torch.nn as nn
import torch.nn.functional as F

class PytorchGridSample(torch.nn.Module):
    def __init__(self, grids):
        super(PytorchGridSample, self).__init__()
        self.upsample1 = nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1)
        self.upsample2 = nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1)
        self.upsample3 = nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1)
        self.upsample4 = nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1)
        self.upsample5 = nn.ConvTranspose2d(32, 3, kernel_size=4, stride=2, padding=1)
        self.grids = grids

    def forward(self, x):
        x = self.upsample1(x)
        x = F.grid_sample(x, self.grids[0], padding_mode='reflection', align_corners=False)
        x = self.upsample2(x)
        x = F.grid_sample(x, self.grids[1], padding_mode='reflection', align_corners=False)
        x = self.upsample3(x)
        x = F.grid_sample(x, self.grids[2], padding_mode='reflection', align_corners=False)
        x = self.upsample4(x)
        x = F.grid_sample(x, self.grids[3], padding_mode='reflection', align_corners=False)
        x = self.upsample5(x)
        x = F.grid_sample(x, self.grids[4], padding_mode='reflection', align_corners=False)
        return x

def convert_to_coreml(model, input_):
    traced_model = torch.jit.trace(model, example_inputs=input_, strict=False)
    coreml_model = ct.converters.convert(
        traced_model,
        inputs=[ct.TensorType(shape=input_.shape)],
        compute_precision=ct.precision.FLOAT16,
        minimum_deployment_target=ct.target.macOS14,
        compute_units=ct.ComputeUnit.ALL
    )
    return coreml_model

def main(pt_model, input_):
    coreml_model = convert_to_coreml(pt_model, input_)
    coreml_model.save("grid_sample.mlpackage")

if __name__ == "__main__":
    input_tensor = torch.randn(1, 512, 4, 4)
    grids = [torch.randn(1, 2*i, 2*i, 2) for i in [4, 8, 16, 32, 64, 128]]
    pt_model = PytorchGridSample(grids)
    main(pt_model, input_tensor)

Bug Report: macOS 15 Beta - PyTorch gridsample Not Utilising Apple Neural Engine on MacBook Pro M2
 
 
Q