tools/autograd/templates/Functions.cpp
// See NOTE [ as_strided Backward ] for explanation
See NOTE [ as_strided Backward ] for explanation
See NOTE [ as_strided Backward ] for explanation
Tensor as_strided_backward(Tensor grad, TensorGeometry input_geometry, IntList sizes, IntList strides, int64_t storage_offset) {
  auto odim = grad.dim();
  std::vector<int64_t> out_sizes_, out_strides_;
  out_sizes_.reserve(odim);
  out_strides_.reserve(odim);
  for (int64_t i = odim - 1; i >= 0; i--) {
    auto size_i = sizes[i];
    auto stride_i = strides[i];
    if (size_i == 0) {
      return at::zeros(input_geometry.sizes(), grad.type());
    } else if (size_i == 1) {
      grad = grad.squeeze(i);
    } else if (stride_i == 0) {
      grad = grad.sum(i, false);
    } else {
      out_sizes_.insert(out_sizes_.begin(), size_i);
      out_strides_.insert(out_strides_.begin(), stride_i);
    }
  }
  auto out_maybe_overlap = _maybe_overlapping_memory(out_sizes_, out_strides_);
  auto idim = input_geometry.dim();
  IntList inp_sizes = input_geometry.sizes(), inp_strides = input_geometry.strides();
  std::vector<int64_t> inp_sizes_, inp_strides_;
  inp_sizes_.reserve(idim);
  inp_strides_.reserve(idim);
  for (int64_t i = idim - 1; i >= 0; i--) {
    auto size_i = inp_sizes[i];
    auto stride_i = inp_strides[i];
    if (size_i == 0) {
      return at::zeros(input_geometry.sizes(), grad.type());
    } else if (size_i != 1) {
      inp_sizes_.insert(inp_sizes_.begin(), size_i);
      inp_strides_.insert(inp_strides_.begin(), stride_i);
    }
  }
  auto inp_maybe_overlap = _maybe_overlapping_memory(inp_sizes_, inp_strides_);
  auto shared_offset = std::min(input_geometry.storage_offset(), storage_offset);
  auto inp_effective_offset = input_geometry.storage_offset() - shared_offset;
  auto out_effective_offset = storage_offset - shared_offset;
  auto base_size = std::max(
    _min_storage_size(inp_sizes_, inp_strides_, inp_effective_offset),
    _min_storage_size(out_sizes_, out_strides_, out_effective_offset)
  );
  auto& ty = grad.type();
  auto storage = at::zeros({base_size}, ty);
  at::optional<at::Tensor> flatten_full_indices;
  if (inp_maybe_overlap || out_maybe_overlap) {
    flatten_full_indices = at::arange(0, base_size, ty.toScalarType(at::kLong));
  }
  if (out_maybe_overlap) {
    auto out_indices = flatten_full_indices->as_strided(out_sizes_, out_strides_, out_effective_offset);
    storage.index_add_(0, out_indices.reshape(-1), grad.reshape(-1));
  } else {
    storage.as_strided(out_sizes_, out_strides_, out_effective_offset).copy_(grad);
  }
  if (inp_maybe_overlap) {
    auto count = at::zeros_like(storage);
    auto inp_indices = flatten_full_indices->as_strided(inp_sizes_, inp_strides_, inp_effective_offset).reshape(-1);
    count.index_add_(0, inp_indices, at::ones({1}, ty).expand_as(inp_indices));
    storage.div_(count); 
  }
  return storage.as_strided(inp_sizes, inp_strides, inp_effective_offset);
}
Tensor as_strided_backward ( Tensor grad , TensorGeometry input_geometry , IntList sizes , IntList strides , int64_t storage_offset ) { auto odim = grad . dim ( ) ; std :: vector < int64_t > out_sizes_ , out_strides_ ; out_sizes_ . reserve ( odim ) ; out_strides_ . reserve ( odim ) ; for ( int64_t i = odim - 1 ; i >= 0 ; i -- ) { auto size_i = sizes [ i ] ; auto stride_i = strides [ i ] ; if ( size_i == 0 ) { return at :: zeros ( input_geometry . sizes ( ) , grad . type ( ) ) ; } else if ( size_i == 1 ) { grad = grad . squeeze ( i ) ; } else if ( stride_i == 0 ) { grad = grad . sum ( i , false ) ; } else { out_sizes_ . insert ( out_sizes_ . begin ( ) , size_i ) ; out_strides_ . insert ( out_strides_ . begin ( ) , stride_i ) ; } } auto out_maybe_overlap = _maybe_overlapping_memory ( out_sizes_ , out_strides_ ) ; auto idim = input_geometry . dim ( ) ; IntList inp_sizes = input_geometry . sizes ( ) , inp_strides = input_geometry . strides ( ) ; std :: vector < int64_t > inp_sizes_ , inp_strides_ ; inp_sizes_ . reserve ( idim ) ; inp_strides_ . reserve ( idim ) ; for ( int64_t i = idim - 1 ; i >= 0 ; i -- ) { auto size_i = inp_sizes [ i ] ; auto stride_i = inp_strides [ i ] ; if ( size_i == 0 ) { return at :: zeros ( input_geometry . sizes ( ) , grad . type ( ) ) ; } else if ( size_i != 1 ) { inp_sizes_ . insert ( inp_sizes_ . begin ( ) , size_i ) ; inp_strides_ . insert ( inp_strides_ . begin ( ) , stride_i ) ; } } auto inp_maybe_overlap = _maybe_overlapping_memory ( inp_sizes_ , inp_strides_ ) ; auto shared_offset = std :: min ( input_geometry . storage_offset ( ) , storage_offset ) ; auto inp_effective_offset = input_geometry . storage_offset ( ) - shared_offset ; auto out_effective_offset = storage_offset - shared_offset ; auto base_size = std :: max ( _min_storage_size ( inp_sizes_ , inp_strides_ , inp_effective_offset ) , _min_storage_size ( out_sizes_ , out_strides_ , out_effective_offset ) ) ; auto & ty = grad . type ( ) ; auto storage = at :: zeros ( { base_size } , ty ) ; at :: optional < at :: Tensor > flatten_full_indices ; if ( inp_maybe_overlap || out_maybe_overlap ) { flatten_full_indices = at :: arange ( 0 , base_size , ty . toScalarType ( at :: kLong ) ) ; } if ( out_maybe_overlap ) { auto out_indices = flatten_full_indices -> as_strided ( out_sizes_ , out_strides_ , out_effective_offset ) ; storage . index_add_ ( 0 , out_indices . reshape ( -1 ) , grad . reshape ( -1 ) ) ; } else { storage . as_strided ( out_sizes_ , out_strides_ , out_effective_offset ) . copy_ ( grad ) ; } if ( inp_maybe_overlap ) { auto count = at :: zeros_like ( storage ) ; auto inp_indices = flatten_full_indices -> as_strided ( inp_sizes_ , inp_strides_ , inp_effective_offset ) . reshape ( -1 ) ; count . index_add_ ( 0 , inp_indices , at :: ones ( { 1 } , ty ) . expand_as ( inp_indices ) ) ; storage . div_ ( count ) ; } return storage . as_strided ( inp_sizes , inp_strides , inp_effective_offset ) ; }
See NOTE [ as_strided Backward ] for explanation
See NOTE [ as_strided Backward ] for explanation
// For output geometry, //   check for size 0 dimensions, //   skip size 1 dimensions, //   reduce grad on expanded dims (stride=0, size>1) // Step (0)     for the algorithm in NOTE [ as_strided Backward ] // Step (0)~(1) for the algorithm in NOTE [ Detecting Memory Overlap Within A Strided Tensor ] //              on output geometry // Step (2)~(4) for the algorithm in NOTE [ Detecting Memory Overlap Within A Strided Tensor ] //              on output geometry // For input geometry, //   check for size 0 dimensions, //   skip size 1 dimensions, // Step (0)~(1) for the algorithm in NOTE [ Detecting Memory Overlap Within A Strided Tensor ] //              on input geometry // Step (1)~(4) for the algorithm in NOTE [ Detecting Memory Overlap Within A Strided Tensor ] //              on input geometry // Rest of this function implements // Step (1)~(4) for the algorithm in NOTE [ as_strided Backward ] // TODO: Raise if not all output values are visible in input geometry. //       Technically speaking, if you treat those values as constants, not //       raising is fine, and mathematically correct. However, these values //       really are contained in some base tensor, and by treating them as //       constants we are ignoring this tight dependency. Therefore, it is //       more sensible to raise here. // Step (1): create underlying tensor as "storage" // prepare indices tensor if we will do index_add_ later // Step (2): use output geometry to scatter gradients into storage // assume that new tensors have 0 storage offset // Step (3): if input tensor has overlapping memory, divide scattered gradient //           at storage[i] by the number of times i shows up in input geometry // this will give nan outside visible range // Step (4): return as_strided view of the storage tensor with input geometry