1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
|
#pragma once
#include <torch/csrc/Export.h>
#include <torch/csrc/jit/ir/ir.h>
#include <memory>
namespace torch::jit {
// Moved from shape_analysis.cpp
// Requirements:
// dims : preserved from the first argument
// scalar type : preserved from the first argument (doesn't have to
// match other arguments)
// device : always matching and preserved
// tensor inputs : *
// tensor outputs : 1
// NB: those ops (with slight adjustments) are good candidates for restarts.
// Knowing the type and device of weights or biases is usually enough to
// infer the output type.
std::shared_ptr<OperatorSet> nn_ops_first_input_preserving();
// Requirements:
// dims : Changed from first argument
// scalar type : preserved from the first argument
// device : always matching and preserved
// tensor inputs : 1
// tensor outputs : 1
std::shared_ptr<OperatorSet> ops_one_tensor_in_shape_transform();
} // namespace torch::jit
|
