tenzr::txpr::VmappedFunction¶

Callable wrapper produced by Vmap that traces the wrapped function unbatched and vectorizes it over a batch dimension. More...

#include <vmap.h>

Public Functions¶

Detailed Description¶

template <typename Fn >
class tenzr::txpr::VmappedFunction;

Callable wrapper produced by Vmap that traces the wrapped function unbatched and vectorizes it over a batch dimension.

Public Functions Documentation¶

function VmappedFunction¶

inline explicit VmappedFunction(
    Fn fn,
    std::vector< int > in_axes ={0},
    int out_axis =0
)

function operator()¶

template <typename... Args>
inline Array operator()(
    const Args &... args
)