Searched defs:iota_r2 (Results 1 – 2 of 2) sorted by relevance
73 xla::XlaOp iota_r2 = xla::Broadcast(iota_r1, {batch_size}); in Compile() local
153 xla::XlaOp iota_r2 = xla::Broadcast(iota_r1, {height}); in CreateTopKUnique() local289 xla::XlaOp iota_r2 = xla::Broadcast(iota_r1, {height}); in CreateMakeUnique() local