microsoft / tensorflow-directml-plugin

DirectML PluggableDevice plugin for TensorFlow 2
Apache License 2.0
185 stars 25 forks source link

Add unsorted segment ops (emulated on the CPU) #253

Closed PatriceVignola closed 2 years ago

PatriceVignola commented 2 years ago

Until we have a DML implementation that doesn't explode in memory usage, we have to emulate this operator on the CPU. We could also simply decide to not implement it, but it has been seen in many models in the wild and breaks colocation graphs.