Since Switch is quite literally the scalar version of IfElse—at least in terms of its Op.c_code implementation, we should consider making invocations of the ifelse helper function return Switches when its arguments are ScalarTypeed.
This should increase the chances that the resulting Op can/will be fused with other scalar Ops to produce a single, more efficient Composite Op, as well as remove some minor overhead due to TensorFromScalar and ScalarFromTensor.
The same general idea can be applied through rewrites, although we'll probably get a lot of coverage using only changes like this at the interface-level. In cases where TensorFromScalar arguments are passed to ifelse, an undesirable IfElse Op would be used, and this is where a rewrite could be helpful.
Originally posted by @brandonwillard in #1217 (comment)
Since
Switchis quite literally the scalar version ofIfElse—at least in terms of itsOp.c_codeimplementation, we should consider making invocations of theifelsehelper function returnSwitches when its arguments areScalarTypeed.This should increase the chances that the resulting
Opcan/will be fused with other scalarOps to produce a single, more efficientCompositeOp, as well as remove some minor overhead due toTensorFromScalarandScalarFromTensor.The same general idea can be applied through rewrites, although we'll probably get a lot of coverage using only changes like this at the interface-level. In cases where
TensorFromScalararguments are passed toifelse, an undesirableIfElseOpwould be used, and this is where a rewrite could be helpful.Originally posted by @brandonwillard in #1217 (comment)