microsoft / tensorflow-directml-plugin

DirectML PluggableDevice plugin for TensorFlow 2
Apache License 2.0
179 stars 23 forks source link

Add int64 support for RandomUniformInt #283

Closed PatriceVignola closed 1 year ago

PatriceVignola commented 1 year ago

int64 support for RandomUniformInt is tricky because DirectML doesn't support the modulus operator for int64 yet, and modulus is necessary to keep the generated integers within the range specified by the user. To work around this, when the range is smaller than uint32_max, we use a special algorithm that allows us to do a modulus of an int64 number by an int32 by doing it separately on its low and high parts. On the other hand, when the range is too big, we fall back to the CPU.