pytorch/torch/csrc/jit/ir/irparser.h
Elias Ellison 43b56b3814 Add Parsing of tensor constants (#75119)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/75119

Add support for parsing Tensor constants like Double(4, 4) ... by initializing random tensors. This makes saving IR and then parsing it lossy, so I have it toggled as default not on, but is useful in cases like repro-ing Fusions with tensor constants post-freezing.

cc Krovatkin

Test Plan: Imported from OSS

Reviewed By: ejguan

Differential Revision: D35373999

Pulled By: eellison

fbshipit-source-id: a5c8d9f93f23a7442258fc745ed6b6def330dca8
(cherry picked from commit 32dd6567522973563bd452bf486ed27b02e4e35c)
2022-04-06 18:00:53 +00:00

41 lines
1.1 KiB
C++

#pragma once
#include <torch/csrc/Export.h>
#include <string>
#include <unordered_map>
#include <c10/util/Optional.h>
#include <torch/csrc/Export.h>
namespace torch {
namespace jit {
struct Graph;
struct Value;
// \brief Parse IR from \p STR constructing the corresponding IR in\ GRAPH.
// if parse_tensor_constants is true will construct empty tensors
// for Tensor constants with random or unitialized contents, otherwise will
// throw
TORCH_API void parseIR(
const std::string& str,
torch::jit::Graph* graph,
bool parse_tensor_constants = false);
/** \brief Parse IR from \p STR constructing the corresponding IR in\ GRAPH.
*
* \p VMAP is filled with String to Value pairs allowing to index Values in the
* newly created graph by their name in the original IR string.
* if parse_tensor_constants is true will construct empty tensors
* for Tensor constants with random or unitialized contents, otherwise will
* throw
*/
TORCH_API void parseIR(
const std::string& str,
torch::jit::Graph* graph,
std::unordered_map<std::string, Value*>& vmap,
bool parse_tensor_constants = false);
} // namespace jit
} // namespace torch