erase_number_types.h

#pragma once

#include <torch/csrc/jit/ir.h>

namespace torch {
namespace jit {

// Erase NumberType information. This is necessary for and only used in
// exporting to ONNX. This pass ensures that no remaining Values have
// NumberType types, replacing them with tensors.
// The following things are done to erase NumberType info:
// - NumberType outputs are changed to DynamicType.
// - prim::Constant nodes which are numbers get changed into 0-dim tensors of
//   the corresponding type
// - prim::TensorToNum, prim::ImplicitTensorToNum and prim::NumToTensor nodes
//   are erased.
//
// The pass assumes that DCE will be called sometime after.
TORCH_API void EraseNumberTypes(const std::shared_ptr<Graph>& graph);

} // namespace jit
} // namespace torch