fn
ge_mask_storage
→Storageint ge_mask_storage(const int & a, const int & b, int numel, Dtype dt, Device device)Element-wise greater-or-equal mask: (a >= b) ? 1 : 0.
Parameters
a, bconst Storage&Inputs to compare.
numelstd::size_tNumber of elements.
dtDtypeOutput dtype the mask is cast to.
deviceDeviceTarget device.
Returns
StorageNew storage holding the gated mask.
Notes
Used by max / clamp-style backward passes to route the gradient only through the "winning" element of each comparison.