pymoose.mux#

pymoose.mux(selector, x, y, placement=None)[source]#

Multiplex two tensors according to some condition.

This op allows for static control-flow of elements coming from two tensors. For boolean tensor s and arbitrary tensors x and y, the operation is equivalent to s * (x - y) + y.

Parameters

  • selector – Boolean tensor representing the control-flow condition.

  • x – A tensor to fill from when the selector condition is 1.

  • y – A tensor to fill from when the selector condition is 0.

  • placement – An optional Placement to pin this operation to.

Result:

A tensor with elements of x and y multiplexed according to the condition given by selector.