Add example usage

python/torch_mlir/csrc/base_lazy_backend/utils/tensor_utils.h

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "torch/csrc/lazy/backend/backend_device.h"
+#include "torch/csrc/lazy/core/tensor.h"
+
+#include "../ops/device_data.h"
+
+
+namespace torch {
+namespace lazy {
+
+inline torch::lazy::DeviceData* device_data_cast(
+    const at::Tensor& tensor, c10::optional<torch::lazy::BackendDevice> device = c10::nullopt
+) {
+    if (!device) {
+        device = torch::lazy::GetBackendDevice(tensor);
+    }
+    TORCH_CHECK(device);
+    torch::lazy::LazyTensorPtr lazy_tensor = torch::lazy::GetLtcTensorOrCreateForWrappedNumber(tensor, *device);
+    if (lazy_tensor) {
+        torch::lazy::Value param_value = lazy_tensor->GetIrValue();
+        if (param_value && param_value->op() == torch::lazy::DeviceData::ClassOpKind()) {
+            return dynamic_cast<torch::lazy::DeviceData*>(param_value.node.get());
+        }
+    }
+    return nullptr;
+}
+
+}  // namespace lazy
+}  // namespace torch

python/torch_mlir/csrc/reference_lazy_backend/reference_lazy_backend_pybind.cpp

@@ -12,6 +12,7 @@
 
 #include <torch_mlir/csrc/base_lazy_backend/mlir_lowering_context.h>
 #include <torch_mlir/csrc/base_lazy_backend/utils/sys_utils.h>
+#include <torch_mlir/csrc/base_lazy_backend/utils/tensor_utils.h>
 
 #include <exception>
 #include <iostream>
@@ -73,6 +74,15 @@ PYBIND11_MODULE(_REFERENCE_LAZY_BACKEND, m) {
         torch::lazy::GetLatestComputation().get());
     return py::cast(computation);
   });
+  m.def("set_parameter_name",
+        [](const at::Tensor& tensor, const std::string& name) -> bool {
+            torch::lazy::DeviceData* ir_node = torch::lazy::device_data_cast(tensor);
+            if (ir_node) {
+                ir_node->SetName(name);
+                return true;
+            }
+            return false;
+        });
   m.def("_initialize", []() {
     NoGilSection gil;
     Initialize();