27#include "neml2/models/map_types.h"
Definition DiagnosticsInterface.cxx:30
c10::Device Device
Definition types.h:66
ValueMap valuemap_no_operation(ValueMap &&x)
No operation.
Definition valuemap_helpers.cxx:58
std::map< LabeledAxisAccessor, Tensor > ValueMap
Definition map_types_fwd.h:33
int64_t Size
Definition types.h:69
ValueMap valuemap_move_device(ValueMap &&x, Device device)
Move all tensors in a ValueMap to a device.
Definition valuemap_helpers.cxx:49
ValueMap valuemap_cat_reduce(std::vector< ValueMap > &&results, Size batch_dim)
Concatenate the tensors in the ValueMap along the batch dimension.
Definition valuemap_helpers.cxx:32