topos 0.12.0

An autodiff compiler stack in Rust.
Documentation
use std::marker::PhantomData;

use static_assertions::assert_impl_all;

use crate::{Element, Shape, Symbol, Tape, Tensor, Value};

use super::Module;

// Entry-time thread-safety contract; the anchor rationale is documented
// in `network.rs`.
assert_impl_all!(Dropout<f64>: Send, Sync);

/// A mask-fed dropout: the expression multiplies its input by a
/// declared mask input whose default payload is all ones, so an unfed
/// run is the identity — inference is the absence of a feed, not a
/// mode.
///
/// Randomness stays outside the graph: masks are generated host-side
/// by the seeded [`init::dropout`](super::init::dropout) factory
/// (inverted dropout, each element `0` or `1 / keep`) and fed per
/// training step like any other input. That keeps seeded runs
/// bit-identical, gradients exact (`d input = gradient * mask`, and
/// the mask edge carries no gradient of its own), and the emitted
/// form of a training step just one more dynamic argument.
///
/// The module holds only the mask input's [`Symbol`]; the keep
/// probability is caller territory, chosen where the mask is drawn.
/// It carries no parameters, so [`Module::visit`] keeps its stateless
/// default.
#[derive(Debug, Clone)]
pub struct Dropout<E> {
    mask: Symbol,
    _marker: PhantomData<E>,
}

impl<E: Element> Dropout<E> {
    /// Declares the mask input on `tape`, shaped like the values
    /// the module will express over, with the all-ones identity
    /// default.
    pub fn new(tape: &Tape<E>, shape: impl Into<Shape>) -> Self {
        let mask = tape.input(Tensor::counted(shape.into(), 1));
        Self {
            mask: mask.symbol(),
            _marker: PhantomData,
        }
    }

    /// Returns the mask input's symbol, for the training loop's feed
    /// pairs.
    pub fn mask(&self) -> Symbol {
        self.mask
    }
}

impl<E: Element> Module<E> for Dropout<E> {
    /// Records the masked value: `input * mask`.
    ///
    /// # Panics
    /// Panics if `input`'s shape differs from the declared mask
    /// shape, or if the module's mask does not resolve on `tape`
    /// generation.
    fn express<'tape>(&self, input: Value<'tape, E>) -> Value<'tape, E> {
        let tape = input.tape();
        input * tape.resolve(self.mask)
    }
}

#[cfg(test)]
#[path = "tests/dropout_tests.rs"]
mod tests;