pub struct WhereRule;Expand description
Backward rule for Op::Where (ternary select).
Forward: out[i] = if cond[i] != 0 { a[i] } else { b[i] }.
Inputs: (cond, a, b) where cond is U8, a and b share
dtype T.
Gradients:
cond:None(U8mask is non-differentiable).a:upstream * cast(cond, T)— gradient flows only at the slots whereawas picked.b:upstream * cast(1 - cond, T)— gradient flows only at the slots wherebwas picked.
The 1 - cond mask is built via Op::AddScalar(-1.0) followed by
Op::Neg: m_b = -(m_a - 1) = 1 - m_a. This avoids needing a
Const(1) factory inside the rule (which would require synthesizing
a slot-populated leaf via the device handle).
Trait Implementations§
Source§impl GradientRule for WhereRule
impl GradientRule for WhereRule
Source§fn backward(
&self,
graph: &SharedGraph,
_op: &Op,
inputs: &[NodeId],
_output: NodeId,
upstream: NodeId,
) -> GradientList
fn backward( &self, graph: &SharedGraph, _op: &Op, inputs: &[NodeId], _output: NodeId, upstream: NodeId, ) -> GradientList
Build gradient nodes for this op. Read more
Auto Trait Implementations§
impl Freeze for WhereRule
impl RefUnwindSafe for WhereRule
impl Send for WhereRule
impl Sync for WhereRule
impl Unpin for WhereRule
impl UnsafeUnpin for WhereRule
impl UnwindSafe for WhereRule
Blanket Implementations§
Source§impl<T> BorrowMut<T> for Twhere
T: ?Sized,
impl<T> BorrowMut<T> for Twhere
T: ?Sized,
Source§fn borrow_mut(&mut self) -> &mut T
fn borrow_mut(&mut self) -> &mut T
Mutably borrows from an owned value. Read more