1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
|
// Copyright (c) Facebook, Inc. and its affiliates.
// All rights reserved.
//
// This source code is licensed under the BSD-style license found in the
// LICENSE file in the root directory of this source tree.
#include <torch/extension.h>
#include <functorch/csrc/dim/dim.h>
namespace at {
namespace functorch {
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
// initialize first-class dims and install it as a submodule on _C
auto dim = Dim_init();
if (!dim) {
throw py::error_already_set();
}
py::setattr(m, "dim", py::reinterpret_steal<py::object>(dim));
}
}}
|
