<commit_before>#pragma once

#include <functional>
#include <initializer_list>
#include <unordered_map>
#include <vector>

#include "xchainer/array.h"
#include "xchainer/constant.h"
#include "xchainer/dtype.h"
#include "xchainer/graph.h"
#include "xchainer/shape.h"

namespace xchainer {

enum class DoubleBackpropOption : bool {
    kDisable = false,
    kEnable = true,
};

class BackwardContext;

namespace backward_detail {

using BackwardFunc = std::function<void(BackwardContext&)>;

// An opaque object class returned by BackwardContext::input_grad().
class SetInputGradProxy {
public:
    explicit SetInputGradProxy(BackwardContext& bctx) : bctx_{bctx} {}

    SetInputGradProxy(SetInputGradProxy&&) = default;
    SetInputGradProxy(const SetInputGradProxy&) = default;
    SetInputGradProxy& operator=(SetInputGradProxy&&) = delete;
    SetInputGradProxy& operator=(const SetInputGradProxy&) = delete;

    // Assign the input gradient.
    template <typename Arg>
    void operator=(Arg&& arg);

private:
    BackwardContext& bctx_;
};

}  // namespace backward_detail

class OpNode;

namespace internal {
class ArrayBody;
}  // namespace internal

class BackwardContext {
public:
    // Ctor
    //
    // `input_grads_storage` is where input gradients returned by backward functions will be stored.
    // Its size must be equal to the number of input arrays whose gradients are to be returned in this single backward function (1 in most
    // ordinary functions).
    BackwardContext(
            const OpNode& op_node,
            Device& output_device,
            gsl::span<const std::reference_wrapper<const nonstd::optional<Array>>> prev_grads,
            gsl::span<const Shape> prev_node_shapes,
            gsl::span<const Dtype> prev_node_dtypes,
            gsl::span<const GraphId> stop_graph_ids,
            gsl::span<std::reference_wrapper<nonstd::optional<Array>>> input_grads_storage);

    // Returns whether the output has a propagated gradient.
    // If there is only one output, the output always has the propagated gradient, therefore you do not have to call this function in that
    // case.
    bool HasOutputGrad(int output_index) const { return gsl::at(prev_grads_, output_index).get().has_value(); }

    // Returns the reference to an output gradient array if it has a propagated value.
    // Otherwise, an zero-filled array is allocated and a reference to it is returned.
    const Array& output_grad(int output_index) const { return GetOutputGrad(output_index); }

    // Returns the reference to an output gradient array if it has a propagated value.
    // Otherwise, an zero-filled array is allocated and a reference to it is returned.
    const Array& output_grad() const {
        assert(prev_grads_.size() == 1);
        return GetOutputGrad(0);
    }

    // Stores the computed input gradient.
    backward_detail::SetInputGradProxy input_grad() { return backward_detail::SetInputGradProxy{*this}; }

    Array Cut(const Array& a) const {
#ifndef NDEBUG
        for (const nonstd::optional<Array>& prev_grad : prev_grads_) {
            assert((!prev_grad.has_value() || &(*prev_grad) != &a) && "Output grads do not have to be cut");
        }
#endif /*NDEBUG*/
        return a.AsConstant(stop_graph_ids_);
    }

private:
    friend class backward_detail::SetInputGradProxy;

    // Stores the computed input gradient.
    void SetInputGrad(Array&& input_grad) {
        assert(input_grads_storage_.size() == 1);
        SetInputGradImpl(input_grads_storage_[0].get(), std::move(input_grad));
    }

    // Stores the computed input gradient.
    void SetInputGrad(const Array& input_grad) {
        assert(input_grads_storage_.size() == 1);
        SetInputGradImpl(input_grads_storage_[0].get(), input_grad);
    }

    // Stores the computed input gradient.
    void SetInputGrad(gsl::span<Array> input_grads) {
        assert(input_grads_storage_.size() == input_grads.size());
        auto it_dst = input_grads_storage_.begin();
        for (const Array& input_grad : input_grads) {
            SetInputGradImpl(it_dst->get(), std::move(input_grad));
            ++it_dst;
        }
    }

    // Stores the computed input gradient.
    void SetInputGrad(std::initializer_list<Array> input_grads) {
        assert(input_grads_storage_.size() == input_grads.size());
        auto it_dst = input_grads_storage_.begin();
        for (const Array& input_grad : input_grads) {
            SetInputGradImpl(it_dst->get(), input_grad);
            ++it_dst;
        }
    }

    void SetInputGradImpl(nonstd::optional<Array>& grad_storage, Array&& input_grad) {
        if (grad_storage.has_value()) {
            grad_storage = *grad_storage + std::move(input_grad);
        } else {
            grad_storage = std::move(input_grad);
        }
    }

    void SetInputGradImpl(nonstd::optional<Array>& grad_storage, const Array& input_grad) {
        if (grad_storage.has_value()) {
            grad_storage = *grad_storage + input_grad;
        } else {
            grad_storage = input_grad;
        }
    }

    // Returns the reference to an output gradient array if it has a propagated value.
    // Otherwise, an zero-filled array is allocated and a reference to it is returned.
    const Array& GetOutputGrad(int output_index) const;

    // Returns the reference to an output gradient array if it has a propagated value.
    // Otherwise, an zero-filled array is allocated and a reference to it is returned.
    const Array& GetOutputGrad() const {
        assert(prev_grads_.size() == 1);
        return GetOutputGrad(0);
    }

    const OpNode& op_node_;
    Device& output_device_;
    gsl::span<const std::reference_wrapper<const nonstd::optional<Array>>> prev_grads_;
    gsl::span<const Shape> prev_node_shapes_;
    gsl::span<const Dtype> prev_node_dtypes_;
    gsl::span<const GraphId> stop_graph_ids_;

    // Gradient passed in SetInputGrad() will be put into this storage.
    gsl::span<std::reference_wrapper<nonstd::optional<Array>>> input_grads_storage_;

    // Holds zero-filled arrays for outputs without actual gradients.
    // The arrays are allocated on-demand in GetOutputGrad.
    mutable std::vector<nonstd::optional<Array>> zero_output_grads_;
};

namespace backward_detail {

template <typename Arg>
void SetInputGradProxy::operator=(Arg&& arg) {
    bctx_.SetInputGrad(std::forward<Arg>(arg));
}

}  // namespace backward_detail

class DefineBackwardScope {
public:
    DefineBackwardScope(const char* op_name, std::initializer_list<ConstArrayRef> outputs, gsl::span<const GraphId> stop_graph_ids);
    DefineBackwardScope(const char* op_name, std::initializer_list<ConstArrayRef> outputs) : DefineBackwardScope{op_name, outputs, {}} {}
    DefineBackwardScope(const char* op_name, const Array& output, gsl::span<const GraphId> stop_graph_ids)
        : DefineBackwardScope{op_name, {output}, stop_graph_ids} {};
    DefineBackwardScope(const char* op_name, const Array& output) : DefineBackwardScope{op_name, {output}, {}} {};

    // Defines a backward function with respect to specified input arrays.
    // For multi-input ops, usually this function is called for each of independent subsets of input arrays.
    template <typename BackwardFunc><commit_msg>Allow multi-array assignment with input_grads()<commit_after>#pragma once

#include <functional>
#include <initializer_list>
#include <unordered_map>
#include <vector>

#include "xchainer/array.h"
#include "xchainer/constant.h"
#include "xchainer/dtype.h"
#include "xchainer/graph.h"
#include "xchainer/shape.h"

namespace xchainer {

enum class DoubleBackpropOption : bool {
    kDisable = false,
    kEnable = true,
};

class BackwardContext;

namespace backward_detail {

using BackwardFunc = std::function<void(BackwardContext&)>;

// An opaque object class returned by BackwardContext::input_grad().
class SetInputGradProxy {
public:
    explicit SetInputGradProxy(BackwardContext& bctx) : bctx_{bctx} {}

    SetInputGradProxy(SetInputGradProxy&&) = default;
    SetInputGradProxy(const SetInputGradProxy&) = default;
    SetInputGradProxy& operator=(SetInputGradProxy&&) = delete;
    SetInputGradProxy& operator=(const SetInputGradProxy&) = delete;

    // Assign the input gradient.
    void operator=(std::initializer_list<Array> grads);

    // Assign the input gradient.
    template <typename Arg>
    void operator=(Arg&& arg);

private:
    BackwardContext& bctx_;
};

}  // namespace backward_detail

class OpNode;

namespace internal {
class ArrayBody;
}  // namespace internal

class BackwardContext {
public:
    // Ctor
    //
    // `input_grads_storage` is where input gradients returned by backward functions will be stored.
    // Its size must be equal to the number of input arrays whose gradients are to be returned in this single backward function (1 in most
    // ordinary functions).
    BackwardContext(
            const OpNode& op_node,
            Device& output_device,
            gsl::span<const std::reference_wrapper<const nonstd::optional<Array>>> prev_grads,
            gsl::span<const Shape> prev_node_shapes,
            gsl::span<const Dtype> prev_node_dtypes,
            gsl::span<const GraphId> stop_graph_ids,
            gsl::span<std::reference_wrapper<nonstd::optional<Array>>> input_grads_storage);

    // Returns whether the output has a propagated gradient.
    // If there is only one output, the output always has the propagated gradient, therefore you do not have to call this function in that
    // case.
    bool HasOutputGrad(int output_index) const { return gsl::at(prev_grads_, output_index).get().has_value(); }

    // Returns the reference to an output gradient array if it has a propagated value.
    // Otherwise, an zero-filled array is allocated and a reference to it is returned.
    const Array& output_grad(int output_index) const { return GetOutputGrad(output_index); }

    // Returns the reference to an output gradient array if it has a propagated value.
    // Otherwise, an zero-filled array is allocated and a reference to it is returned.
    const Array& output_grad() const {
        assert(prev_grads_.size() == 1);
        return GetOutputGrad(0);
    }

    // Stores the computed input gradient.
    backward_detail::SetInputGradProxy input_grad() { return backward_detail::SetInputGradProxy{*this}; }

    Array Cut(const Array& a) const {
#ifndef NDEBUG
        for (const nonstd::optional<Array>& prev_grad : prev_grads_) {
            assert((!prev_grad.has_value() || &(*prev_grad) != &a) && "Output grads do not have to be cut");
        }
#endif /*NDEBUG*/
        return a.AsConstant(stop_graph_ids_);
    }

    // Stores the computed input gradient.
    void SetInputGrad(std::initializer_list<Array> input_grads) {
        assert(input_grads_storage_.size() == input_grads.size());
        auto it_dst = input_grads_storage_.begin();
        for (const Array& input_grad : input_grads) {
            SetInputGradImpl(it_dst->get(), input_grad);
            ++it_dst;
        }
    }

private:
    friend class backward_detail::SetInputGradProxy;

    // Stores the computed input gradient.
    void SetInputGrad(Array&& input_grad) {
        assert(input_grads_storage_.size() == 1);
        SetInputGradImpl(input_grads_storage_[0].get(), std::move(input_grad));
    }

    // Stores the computed input gradient.
    void SetInputGrad(const Array& input_grad) {
        assert(input_grads_storage_.size() == 1);
        SetInputGradImpl(input_grads_storage_[0].get(), input_grad);
    }

    // Stores the computed input gradient.
    void SetInputGrad(gsl::span<Array> input_grads) {
        assert(input_grads_storage_.size() == input_grads.size());
        auto it_dst = input_grads_storage_.begin();
        for (const Array& input_grad : input_grads) {
            SetInputGradImpl(it_dst->get(), std::move(input_grad));
            ++it_dst;
        }
    }

    void SetInputGradImpl(nonstd::optional<Array>& grad_storage, Array&& input_grad) {
        if (grad_storage.has_value()) {
            grad_storage = *grad_storage + std::move(input_grad);
        } else {
            grad_storage = std::move(input_grad);
        }
    }

    void SetInputGradImpl(nonstd::optional<Array>& grad_storage, const Array& input_grad) {
        if (grad_storage.has_value()) {
            grad_storage = *grad_storage + input_grad;
        } else {
            grad_storage = input_grad;
        }
    }

    // Returns the reference to an output gradient array if it has a propagated value.
    // Otherwise, an zero-filled array is allocated and a reference to it is returned.
    const Array& GetOutputGrad(int output_index) const;

    // Returns the reference to an output gradient array if it has a propagated value.
    // Otherwise, an zero-filled array is allocated and a reference to it is returned.
    const Array& GetOutputGrad() const {
        assert(prev_grads_.size() == 1);
        return GetOutputGrad(0);
    }

    const OpNode& op_node_;
    Device& output_device_;
    gsl::span<const std::reference_wrapper<const nonstd::optional<Array>>> prev_grads_;
    gsl::span<const Shape> prev_node_shapes_;
    gsl::span<const Dtype> prev_node_dtypes_;
    gsl::span<const GraphId> stop_graph_ids_;

    // Gradient passed in SetInputGrad() will be put into this storage.
    gsl::span<std::reference_wrapper<nonstd::optional<Array>>> input_grads_storage_;

    // Holds zero-filled arrays for outputs without actual gradients.
    // The arrays are allocated on-demand in GetOutputGrad.
    mutable std::vector<nonstd::optional<Array>> zero_output_grads_;
};

namespace backward_detail {

inline void SetInputGradProxy::operator=(std::initializer_list<Array> grads) { bctx_.SetInputGrad(grads); }

template <typename Arg>
void SetInputGradProxy::operator=(Arg&& arg) {
    bctx_.SetInputGrad(std::forward<Arg>(arg));
}

}  // namespace backward_detail

class DefineBackwardScope {
public:
    DefineBackwardScope(const char* op_name, std::initializer_list<ConstArrayRef> outputs, gsl::span<const GraphId> stop_graph_ids);
    DefineBackwardScope(const char* op_name, std::initializer_list<ConstArrayRef> outputs) : DefineBackwardScope{op_name, outputs, {}} {}
    DefineBackwardScope(const char* op_name, const Array& output, gsl::span<const GraphId> stop_graph_ids)
        : DefineBackwardScope{op_name, {output}, stop_graph_ids} {};
    DefineBackwardScope(const char* op_name, const Array& output) : DefineBackwardScope{op_name, {output}, {}} {};

    // Defines a backward function with respect to specified input arrays.
    // For multi-input ops, usually this function is called for each of independent subsets of input arrays.
    template <typename BackwardFunc>
{'fraction_non_alphanumeric': 0.10293119826299363, 'fraction_numerical': 0.0012213326095806758, 'mean_word_length': 3.6305372290292177, 'pattern_counts': {'":': 0, '<': 81, '<?xml version=': 0, '>': 85, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2715844', 'n_tokens_mistral': 4324, 'n_tokens_neox': 4242, 'n_words': 1187}