Func Class — pytorch Architecture

Source Code

aten/src/ATen/functorch/BatchRulesFactory.cpp lines 15–30

template <typename F, F Func, typename A, typename B, typename... T>
struct NewBlahBatchRuleHelperSymInt<F, Func, typelist<A, B, T...>> {
  static std::tuple<Tensor, std::optional<int64_t>> apply(
      const Tensor& tensor,
      std::optional<int64_t> batch_dim,
      SymIntArrayRef shape,
      T... extra_args) {
    // NOLINTNEXTLINE(bugprone-unchecked-optional-access)
    const auto bdim_size = tensor.sym_size(batch_dim.value());
    c10::SmallVector<c10::SymInt> new_shape;
    new_shape.reserve(shape.size() + 1);
    new_shape.emplace_back(bdim_size);
    new_shape.insert(new_shape.end(), shape.begin(), shape.end());
    return std::make_tuple(Func(tensor, new_shape, std::forward<T>(extra_args)...), 0);
  }
};