FM / ATen

/usr/local/lib64/python3.6/site-packages/torch/include/ATen/

/usr/local/lib64/python3.6/site-packages/torch/include/ATen

Name	Size	Mode	Actions
core/	-	0755	rm
cpu/	-	0755	rm
cuda/	-	0755	rm
cudnn/	-	0755	rm
detail/	-	0755	rm
hip/	-	0755	rm
native/	-	0755	rm
quantized/	-	0755	rm
AccumulateType.h	4438	0644	edit dl rm
ArrayRef.h	44	0644	edit dl rm
ATen.h	998	0644	edit dl rm
autocast_mode.h	6716	0644	edit dl rm
Backend.h	43	0644	edit dl rm
Backtrace.h	46	0644	edit dl rm
BatchedFallback.h	965	0644	edit dl rm
BatchedTensorImpl.h	5383	0644	edit dl rm
CompositeExplicitAutogradFunctions.h	1622	0644	edit dl rm
CompositeExplicitAutogradFunctions_inl.h	54075	0644	edit dl rm
CompositeImplicitAutogradFunctions.h	1622	0644	edit dl rm
CompositeImplicitAutogradFunctions_inl.h	142082	0644	edit dl rm
Config.h	734	0644	edit dl rm
Context.h	12767	0644	edit dl rm
cpp_custom_type_hack.h	5326	0644	edit dl rm
CPUApplyUtils.h	12582	0644	edit dl rm
CPUFixedAllocator.h	830	0644	edit dl rm
CPUFunctions.h	1600	0644	edit dl rm
CPUFunctions_inl.h	171924	0644	edit dl rm
CPUGeneratorImpl.h	1431	0644	edit dl rm
CUDAFunctions.h	1601	0644	edit dl rm
CUDAFunctions_inl.h	185696	0644	edit dl rm
CUDAGeneratorImpl.h	4695	0644	edit dl rm
Device.h	42	0644	edit dl rm
DeviceGuard.h	1134	0644	edit dl rm
Dimname.h	31	0644	edit dl rm
DimVector.h	46	0644	edit dl rm
Dispatch.h	52137	0644	edit dl rm
div_rtn.h	204	0644	edit dl rm
DLConvertor.h	576	0644	edit dl rm
dlpack.h	5244	0644	edit dl rm
DynamicLibrary.h	369	0644	edit dl rm
ExpandUtils.h	14506	0644	edit dl rm
Formatting.h	34	0644	edit dl rm
Functions.h	846326	0644	edit dl rm
Generator.h	46	0644	edit dl rm
InferSize.h	2143	0644	edit dl rm
InitialTensorOptions.h	445	0644	edit dl rm
Layout.h	42	0644	edit dl rm
MapAllocator.h	2999	0644	edit dl rm
MatrixRef.h	3016	0644	edit dl rm
MemoryOverlap.h	1117	0644	edit dl rm
MetaFunctions.h	1601	0644	edit dl rm
MetaFunctions_inl.h	84006	0644	edit dl rm
NamedTensor.h	35	0644	edit dl rm
NamedTensorUtils.h	5747	0644	edit dl rm
NativeFunctions.h	354651	0644	edit dl rm
NativeMetaFunctions.h	35445	0644	edit dl rm
NumericUtils.h	2787	0644	edit dl rm
OpaqueTensorImpl.h	6080	0644	edit dl rm
Operators.h	1707199	0644	edit dl rm
OpMathType.h	460	0644	edit dl rm
Parallel.h	4875	0644	edit dl rm
ParallelNative.h	2443	0644	edit dl rm
ParallelNativeTBB.h	2934	0644	edit dl rm
ParallelOpenMP.h	3049	0644	edit dl rm
PTThreadPool.h	394	0644	edit dl rm
record_function.h	24044	0644	edit dl rm
RedispatchFunctions.h	1112886	0644	edit dl rm
RegistrationDeclarations.h	545777	0644	edit dl rm
SavedTensorHooks.h	328	0644	edit dl rm
Scalar.h	44	0644	edit dl rm
ScalarOps.h	2272	0644	edit dl rm
ScalarType.h	129	0644	edit dl rm
SequenceNumber.h	373	0644	edit dl rm
SmallVector.h	47	0644	edit dl rm
SparseCsrTensorImpl.h	2045	0644	edit dl rm
SparseCsrTensorUtils.h	523	0644	edit dl rm
SparseTensorImpl.h	12417	0644	edit dl rm
SparseTensorUtils.h	4219	0644	edit dl rm
Storage.h	43	0644	edit dl rm
Tensor.h	48	0644	edit dl rm
TensorAccessor.h	51	0644	edit dl rm
TensorGeometry.h	1855	0644	edit dl rm
TensorIndexing.h	21923	0644	edit dl rm
TensorIterator.h	29962	0644	edit dl rm
TensorIteratorInternal.h	1862	0644	edit dl rm
TensorMeta.h	2917	0644	edit dl rm
TensorNames.h	2519	0644	edit dl rm
TensorOperators.h	3275	0644	edit dl rm
TensorOptions.h	49	0644	edit dl rm
TensorUtils.h	5687	0644	edit dl rm
ThreadLocalState.h	3289	0644	edit dl rm
TracerMode.h	5576	0644	edit dl rm
TypeDefault.h	680	0644	edit dl rm
Utils.h	5993	0644	edit dl rm
Version.h	340	0644	edit dl rm
VmapMode.h	952	0644	edit dl rm
VmapTransforms.h	7654	0644	edit dl rm
WrapDimUtils.h	3438	0644	edit dl rm
WrapDimUtilsMulti.h	768	0644	edit dl rm

Edit: /usr/local/lib64/python3.6/site-packages/torch/include/ATen/VmapTransforms.h (7654B)

#pragma once

#include <ATen/BatchedTensorImpl.h>

namespace at {

// This file contains abstractions used for transforming *logical* vmap arguments
// into *physical* arguments. (Keep reading for definitions of these terms).

// NOTE: [Logical vs physical args]
// Consider the following vmap.
//   vmap(vmap(func, in_dims=(2,)), in_dims=(0,))(torch.ones(2, 3, 4))
// This would produce a BatchedTensor wrapping a Tensor of size [2, 3, 4],
// with batch dims 0 and 2:
//   BatchedTensor(ones(2, 3, 4), bdims=[(lvl=1,dim=0),(lvl=2,dim=2)])
//
// We say the *logical* view of the tensor has size [3] -- tensors inside
// `func` appear to have size [3].
// However, the *physical* underlying tensor (the one passed to vmap) has size
// [2, 3, 4].
//
// This notion of logical vs physical also extends to non-tensor arguments.
// Consider the previous tensor; let's assume the user called
// `torch.sum(tensor, dim=0)` inside of `func`. Then the logical
// dimension they are reducing over is dim 0 but the physical dim is dim 1
// (the first non-batch dimension)

// Forward declared; see NOTE: [What is a VmapPhysicalView?]
struct VmapPhysicalView;

// Most PyTorch operators take 4 or fewer inputs.
constexpr int64_t kVmapTransformStaticInputSize = 4;
using VmapPhysicalViewVec = SmallVector<VmapPhysicalView, kVmapTransformStaticInputSize>;

// Pytorch generally advertises good performance for <= 5 dims.
// (see ATen/core/DimVector.h). We add a few extra dims (~3) for vmap
// dimensions to get 8. Adjust this number as necessary
constexpr int64_t kVmapStaticDimVecSize = 8;
using VmapDimVector = SmallVector<int64_t, kVmapStaticDimVecSize>;

// NOTE: [What is an VmapTransform?]
// An *VmapTransform* converts logical views of tensors to physical views.
//
// Batching rules use VmapTransforms to convert logical arguments to
// physical arguments, then call one or more at:: operator that handles the
// physical arguments, and then converts the physical result back to a logical
// argument.

// VmapTransform for operators that take tensors with multiple batch dims.
// Given one or more logical views on Tensors, `logicalToPhysical`
// permutes all of the batch dims to the front of the tensor, aligns
// and expands the batch dims to match each other (according to their `level`),
// and returns a VmapPhysicalView on the tensor(s).
struct TORCH_API MultiBatchVmapTransform {
  static VmapPhysicalView logicalToPhysical(const Tensor& logical_tensor);
  static VmapPhysicalViewVec logicalToPhysical(TensorList logical_tensors);
};

// VmapTransform for operators that broadcast all inputs.
// Given some logical views on Tensors, `logicalToPhysical`:
// - permutes all of the batch dims to the front of the tensors
// - aligns all the batch dims to the collective levels of all of the tensors.
//   If a tensor does not have a batch dim for a vmap level, then it receives
//   a size-one dimension for said level.
// - aligns the non-batch dims to have the same dimensionality, adding extra
//   size-1 dimensions in between the batch dimensions and the non-batch dimensions
//   so that the batch dimensions are lined up from the right.
//
// For example: given inputs of size (B, 2) and (B, 3, 2) where B is the batch
// dimension, BroadcastingVmapTransform returns VmapPhysicalViews that wrap tensors
// of size (B, 1, 2) and (B, 3, 2).
//
// Given inputs of size (B, 2) and (2,), BroadcastingVmapTransform returns
// VmapPhysicalViews wrapping tensors of size (B, 2) and (1, 2). We don't
// actually *need* to return a tensor of size (1, 2) for the second tensor
// because the broadcasting operation takes care of that for us, but we do
// it anyways to keep things simple.
struct TORCH_API BroadcastingVmapTransform {
  static VmapPhysicalViewVec logicalToPhysical(TensorList logical_tensors);
};

// Forward declared, if you're reading this file head to toe, don't worry about
// it yet.
struct VmapPhysicalToLogicalMap;

// NOTE: [What is a VmapPhysicalView?]
// VmapPhysicalView represents a physical view on a Tensor.
//
// One can use it to further convert logical dimension indices, logical shapes,
// and more to their physical variants, or convert a new (physical) tensor into
// a logical BatchedTensor. (TODO(rzou): some of these are not yet implemented).
//
// VmapPhysicalView stores a physical tensor with all of its batch dimensions at
// the front and some levels that correspond to said batch dimensions.
//
// The levels bitset specifies which vmap levels correspond to the batch
// dimensions at the front of the tensor. In particular, the number of set bits
// corresponds to the number of batch dimensions on `tensor` and the rightmost
// bit of `levels` specifies the maximum number of nested vmaps we are in at
// this point in time.
// For example, given:
//   physical_view = VmapPhysicalView(tensor=ones(2, 3, 4, 5, 6), levels={1, 3})
//
// Rightmost bit of `levels` is 3 indicating the number of nested vmaps less
// than or equal to 3.
//   bitset: 010100
//              ^
//              |
//   levels: 012345
struct TORCH_API VmapPhysicalView {
  VmapPhysicalView(Tensor&& tensor, std::bitset<kVmapNumLevels> levels)
      : levels_(levels), tensor_(tensor) {
    TORCH_INTERNAL_ASSERT(!isBatchedTensor(tensor));
  }

Tensor& tensor() { return tensor_; }
  const Tensor& tensor() const { return tensor_; }

// Maps logical dim indices to physical dim indices. Also does dim wrapping.
  //
  // For example, given:
  //   physical_view = VmapPhysicalView(tensor=ones(2, 3, 4, 5), levels={1, 3})
  //
  // Then physical_view.getPhysicalDims({0, 1}) returns {2, 3}.
  // This is because the size of levels tell us that the first two dimensions
  // of `tensor_` are batch dimensions, so a logical dim of `n` is actually
  // a physical dim of `n + 2`.
  VmapDimVector getPhysicalDims(IntArrayRef logical_dims) const;
  int64_t getPhysicalDim(int64_t logical_dim) const;

// Returns a VmapPhysicalToLogicalMap object. This can be used for
  // mapping a physical tensor to a new logical tensor (BatchedTensor)
  VmapPhysicalToLogicalMap getPhysicalToLogicalMap() const;

// Maps a logical shape to a physical shape by pre-pending the batch
  // sizes to the logical shape.
  VmapDimVector getPhysicalShape(IntArrayRef logical_shape) const;

int64_t numBatchDims() const;

private:
  int64_t numLogicalDims() const;

std::bitset<kVmapNumLevels> levels_;
  Tensor tensor_;
};

// Convenience struct used for mapping a physical tensor (a non-BatchedTensor)
// to a logical one (BatchedTensor). It holds some levels that are used to do the
// mapping and assumes that the batch dimensions in the physical tensor all
// occur at the front of the tensor.
struct TORCH_API VmapPhysicalToLogicalMap {
  VmapPhysicalToLogicalMap(std::bitset<kVmapNumLevels> levels): levels_(levels) {}

// Maps a physical tensor to a new logical tensor (BatchedTensor).
  // Assumes that all of the "batch dimensions" are at the front
  // of the physical tensor. For example, given:
  // - x = rank-4 Tensor with size 2, 3, 5, 7
  // - levels = (2, 4)
  // Returns:
  // - BatchedTensor(x, bdims=[(dim=0,lvl=2), (dim=1, lvl=4)])
  Tensor apply(const Tensor& physical_tensor) const;

// Given a vector of physical tensors,
  // 1. maps each tensor to a new logical tensor. Assumes that all of the
  //    "batch dimensions" are at the front of the physical tensors.
  // 2. stores the new logical tensors back into the passed-in vector. This is
  //    to avoid additional dynamic allocations.
  void applyInplace(std::vector<Tensor>& physical_tensors) const;

std::bitset<kVmapNumLevels> levels_;
};

} // namespace at

cmd: