#pragma once #include #include #include #include #include namespace at { // This class allows you to write variadic functions which // call a (possibly overloaded) function on each argument, // in order. This is most commonly used in autogenerated code, // where it is convenient to have a function that can uniformly // take arguments of different types. If your arguments // are homogenous consider using a std::initializer_list instead. // // For examples of this in use, see torch/csrc/utils/variadic.h template struct IterArgs { template inline F& apply() { return self(); } // NB: Use perfect forwarding here, otherwise we'll make value // copies of all arguments! template inline F& apply(T&& arg, Args&&... args) { self()(std::forward(arg)); if (self().short_circuit()) { return self(); } else { return apply(std::forward(args)...); } } // Here are some handy overloads which provide sensible // defaults for container-like structures that one might // be interested in recursing into. You can enable them // by adding: // // using IterArgs::operator() // // to your struct. These are not enabled by default because // you may be able to process these structures more efficiently // than handling them one-by-one. template void operator()(at::ArrayRef args) { for (const auto& arg : args) { self()(arg); if (self().short_circuit()) return; } } // NB: we need to specify std::vector manually as C++ won't // do an implicit conversion to make a template deduction go through. template void operator()(const std::vector& args) { self()(at::ArrayRef{args}); } constexpr bool short_circuit() const { return false; } private: inline F& self() { return *static_cast(this); } }; } // namespace torch