is_all Class — pytorch Architecture

Source Code

aten/src/ATen/native/ReduceOps.cpp lines 1674–1711

template <bool is_all>
static Tensor allany_dims_default(const Tensor &self, OptionalIntArrayRef dim, bool keepdim) {
  // Default implementation in terms of all-reduce or single dim reduce
  if (!dim) {
    Tensor out;
    if constexpr (is_all) {
      out = self.all();
    } else {
      out = self.any();
    }

    if (keepdim) {
      DimVector out_shape(self.dim(), 1);
      return out.expand(out_shape);
    }
    return out;
  }

  if (dim->empty()) {
    if (self.scalar_type() == kByte) {
      // Convert to a 1 or 0 mask
      auto out = at::empty_like(self);
      return at::ne_outf(self, 0, out);
    } else {
      return at::_to_copy(self, kBool);
    }
  }

  Tensor out = self;
  for (auto d : *dim) {
    if constexpr (is_all) {
      out = out.all(d, /*keepdim=*/true);
    } else {
      out = out.any(d, /*keepdim=*/true);
    }
  }
  return keepdim ? out : out.squeeze(*dim);
}