Func Class — pytorch Architecture

Source Code

aten/src/ATen/functorch/BatchRulesRandomness.cpp lines 21–36

template <typename F, F Func, typename... ExtraArgs>
static Tensor random_batching_rule(SymIntArrayRef shape, ExtraArgs... extra_args) {
  c10::impl::ExcludeDispatchKeyGuard guard(DispatchKey::FuncTorchVmapMode);
  auto maybe_layer = maybeCurrentDynamicLayer();
  TORCH_INTERNAL_ASSERT(maybe_layer.has_value());
  c10::SmallVector<SymInt> shapeVec(1, maybe_layer->batchSize());
  shapeVec.reserve(shape.size() + 1);
  shapeVec.insert(shapeVec.end(), shape.begin(), shape.end());
  RandomnessType randomness = maybe_layer->randomness();
  check_randomness(randomness);
  if (randomness == RandomnessType::Different) {
    return makeBatched(Func(shapeVec, std::forward<ExtraArgs>(extra_args)...), 0, maybe_layer->layerId());
  } else {
    return Func(shape, std::forward<ExtraArgs>(extra_args)...);
  }
}