Add tracing support for optional Device and Layout

Author: izdeby

torch/csrc/jit/tracer.cpp

@@ -504,6 +504,30 @@ void addInputs(
     n->addInput(none);
   }
 }
+void addInputs(
+    Node* n,
+    const char* name,
+    const c10::optional<at::Layout>& value) {
+  if (value.has_value()) {
+    detail::genericAddInput(n, static_cast<int64_t>(*value));
+  } else {
+    Graph* g = n->owningGraph();
+    Value* none = g->insertNode(g->createNone())->output();
+    n->addInput(none);
+  }
+}
+void addInputs(
+    Node* n,
+    const char* name,
+    const c10::optional<at::Device>& value) {
+  if (value.has_value()) {
+    detail::genericAddInput(n, value);
+  } else {
+    Graph* g = n->owningGraph();
+    Value* none = g->insertNode(g->createNone())->output();
+    n->addInput(none);
+  }
+}
 #ifdef BUILD_NAMEDTENSOR
 void addInputs(
     Node* n,

torch/csrc/jit/tracer.h

@@ -259,6 +259,14 @@ TORCH_API void addInputs(
     Node* n,
     const char* name,
     const c10::optional<at::ScalarType>& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const c10::optional<at::Device>& value);
+TORCH_API void addInputs(
+    Node* n,
+    const char* name,
+    const c10::optional<at::Layout>& value);
 TORCH_API void addInputs(Node* n, const char* name, at::MemoryFormat value);
 #ifdef BUILD_NAMEDTENSOR
 TORCH_API void addInputs(Node* n, const char* name, c10::optional<at::DimnameList> value);