search

apple / coremltools

Core ML tools contain supporting tools for Core ML model conversion, editing, and validation.
https://coremltools.readme.io
BSD 3-Clause "New" or "Revised" License
4.32k stars 627 forks source link

torch.nn.functional.grid_sample for volumetric input/ Rank 5 input and grid #2284

Open nighting0le01 opened 1 month ago

nighting0le01 commented 1 month ago

🌱 Describe your Feature Request

support torch.nn.function.grid_sample for 5d Volumetric input by supporting 5d input for mb.affine and mb.resample for rank >=4 :https://pytorch.org/docs/stable/generated/torch.nn.functional.grid_sample.html

How can this feature be used?

Please provide some examples where this feature can be used. For video generation models features on device

Describe alternatives you've considered

currently grid_sampler is supported for rank 4 tensors. it erros out in resample op mb.resample input "x" to the "resample" op must be a rank 4 tensor. Got rank 5 tensor of shape (22, 4, c,h,w)

please support to get mb.affine and mb.resample for rank 5 inputs

nighting0le01 commented 1 month ago

@junpeiz @TobyRoseman @aseemw this feature is needed for video gen application

jakesabathia2 commented 1 month ago

This sounds like a limitation in the Framework :)

nighting0le01 commented 1 month ago

hello @jakesabathia2 what are the steps to support such a custom layer? do we have a up to date example. So i can implement and land the PR