5#ifndef RLPACK_BINARIES_GRAD_ACCUMULATOR_CGRADACCUMULATOR_H_
6#define RLPACK_BINARIES_GRAD_ACCUMULATOR_CGRADACCUMULATOR_H_
8#include <pybind11/pybind11.h>
9#include <torch/extension.h>
25 C_GradAccumulator(std::vector<std::string> ¶meterKeys, int64_t boostStrapRounds);
28 void accumulate(std::map<std::string, torch::Tensor> &namedParameters);
30 std::map<std::string, torch::Tensor>
sum_reduce();
31 std::map<std::string, torch::Tensor>
get_item(int64_t index);
32 void set_item(int64_t index, std::map<std::string, torch::Tensor> &namedParameters);
Definition: C_GradAccumulator.h:23
void clear()
Definition: C_GradAccumulator.cpp:159
size_t size()
Definition: C_GradAccumulator.cpp:155
int64_t bootstrapRounds_
The number of boostrap rounds over which accumulation and reduction is to take place.
Definition: C_GradAccumulator.h:39
std::map< std::string, torch::Tensor > reducedParams_
The map to store final results of reduced parameters.
Definition: C_GradAccumulator.h:45
std::map< std::string, torch::Tensor > mean_reduce()
Definition: C_GradAccumulator.cpp:43
std::map< std::string, torch::Tensor > sum_reduce()
Definition: C_GradAccumulator.cpp:69
void accumulate(std::map< std::string, torch::Tensor > &namedParameters)
Definition: C_GradAccumulator.cpp:21
std::map< std::string, torch::Tensor > get_item(int64_t index)
Definition: C_GradAccumulator.cpp:95
void set_item(int64_t index, std::map< std::string, torch::Tensor > &namedParameters)
Definition: C_GradAccumulator.cpp:113
void delete_item(int64_t index)
Definition: C_GradAccumulator.cpp:139
std::vector< std::map< std::string, torch::Tensor > > namedParametersGrads_
The vector to accumulate the gradients.
Definition: C_GradAccumulator.h:43
std::vector< std::string > parameterKeys_
The parameter keys for the model for which gradient accumulation is being done.
Definition: C_GradAccumulator.h:41