search

google-ai-edge / ai-edge-torch

Supporting PyTorch models with the Google AI Edge TFLite runtime.
Apache License 2.0
310 stars 40 forks source link

[Prototype] Align SDPA definition with the gpu accelerator #78

Open Linchenn opened 3 months ago

Linchenn commented 3 months ago

cc @haozha111 @advaitjain

With this pass, the converted tflite could be run on the GPU accelerator (though it has a tiny bug to be fixed).

If we finally decide to use this approach, we additionally need to figure out how to diverge conversion between the CPU accelerator and the GPU accelerator.

haozha111 commented 3 months ago

cc @haozha111 @advaitjain

With this pass, the converted tflite could be run on the GPU accelerator (though it has a tiny bug to be fixed).

If we finally decide to use this approach, we additionally need to figure out how to diverge conversion between the CPU accelerator and the GPU accelerator.

thanks! can we make this PR in draft mode? we don't need any code reviews for it now.