grad_def.hpp

Go to the documentation of this file.

#include <list>

#include "teq/traveler.hpp"

#ifndef TEQ_GRAD_DEF_HPP
#define TEQ_GRAD_DEF_HPP

namespace teq
{

struct iGradientBuilder
{
    virtual ~iGradientBuilder (void) = default;

    virtual TensptrT local_derivative (FuncptrT op, size_t arg_idx) const = 0;

    virtual TensptrT chain_rule (FuncptrT op, const TensptrT& local_der,
        TensptrT supcomp_grad, size_t arg_idx) const = 0;

    virtual TensptrT get_const_one (Shape shape) const = 0;

    virtual TensptrT get_const_zero (Shape shape) const = 0;

    virtual TensptrT add (TensptrT& lhs, TensptrT& rhs) const = 0;

    TensptrT derive (TensptrT root, TensptrT target) const
    {
        if (root == target)
        {
            return get_const_one(target->shape());
        }

        PathFinder finder(target.get());
        root->accept(finder);

        auto& pathmap = finder.parents_;
        // no path to wrt
        if (pathmap.empty())
        {
            return get_const_zero(target->shape());
        }
        // else there exists a path to wrt
        // using pathfinder, breadth first traverse from this to wrt
        GraphStat stat;
        root->accept(stat);
        auto owners = track_owners({root});

        std::list<iFunctor*> parents;
        std::transform(pathmap.begin(), pathmap.end(),
            std::back_inserter(parents),
            [](std::pair<iTensor*,std::vector<size_t>> parent)
            {
                return static_cast<iFunctor*>(parent.first);
            });
        parents.sort(
            [&](iFunctor* a, iFunctor* b)
            {
                size_t aheight = stat.graphsize_[a].upper_;
                size_t bheight = stat.graphsize_[b].upper_;
                if (aheight == bheight) // make traversal more deterministic
                {
                    return a->to_string() > b->to_string();
                }
                return aheight > bheight;
            });

        // map functor to its respective super composite derivative
        // let L = root, F = key functor, value of F in grads is dL/dF
        std::unordered_map<const iTensor*,TensptrsT> grads = {
            {root.get(), {get_const_one(root->shape())}}
        };
        for (iFunctor* parent : parents)
        {
            TensptrsT& gargs = grads[parent];
            TensptrT bwd = gargs[0];
            for (size_t i = 1, n = gargs.size(); i < n; ++i)
            {
                bwd = add(bwd, gargs[i]);
            }

            auto& grad_indices = pathmap[parent];
            ArgsT children = parent->get_children();
            size_t nchildren = children.size();
            // assert: all nnary-children use identity mapping,
            // so no children-arg is direct mapping
            TensptrsT args(nchildren);
            std::transform(children.begin(), children.end(), args.begin(),
                [](FuncArg& arg)
                {
                    return arg.get_tensor();
                });
            // for each painted child, calculate dThis/dChild
            // go through grads in order
            std::list<size_t> ordered(grad_indices.begin(), grad_indices.end());
            ordered.sort();
            for (size_t i : ordered)
            {
                auto parent_ptr = std::static_pointer_cast<iFunctor>(
                    owners[parent].lock());
                auto local = local_derivative(parent_ptr, i);
                auto grad_step = chain_rule(parent_ptr, local, bwd, i);
                grads[args[i].get()].push_back(grad_step);
            }
        }
        TensptrsT& outargs = grads[target.get()];
        TensptrT out = outargs[0];
        for (size_t i = 1, n = outargs.size(); i < n; ++i)
        {
            out = add(out, outargs[i]);
        }
        return out;
    }
};

}

#endif // TEQ_GRAD_DEF_HPP