PaddlePaddle
diff --git a/‎paddle/fluid/pir/serialize_deserialize/src/save_load_parameters.cc
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions b/‎paddle/fluid/pir/serialize_deserialize/src/save_load_parameters.cc
Copy file name to clipboard
+9Lines changed: 9 additions & 0 deletions
diff --git a/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+23Lines changed: 23 additions & 0 deletions b/‎paddle/fluid/pybind/pir.cc
Copy file name to clipboard
+23Lines changed: 23 additions & 0 deletions
diff --git a/‎python/paddle/jit/api.py
Copy file name to clipboard
+11-7Lines changed: 11 additions & 7 deletions b/‎python/paddle/jit/api.py
Copy file name to clipboard
+11-7Lines changed: 11 additions & 7 deletions
@@ -25,6 +25,15 @@ const phi::DeviceContext* GetDeviceContext(const phi::DenseTensor& x) {
   phi::DeviceContextPool& pool = phi::DeviceContextPool::Instance();
   const phi::DeviceContext* dev_ctx = nullptr;
   auto place = x.place();
+  if (place.GetType() == phi::AllocationType::UNDEFINED) {
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
+    place = phi::GPUPlace();
+#elif defined(PADDLE_WITH_XPU)
+    place = phi::XPUPlace();
+#else
+    place = phi::CPUPlace();
+#endif
+  }
   dev_ctx = pool.Get(place);
   return dev_ctx;
 }
 
@@ -897,12 +897,35 @@ void BindValue(py::module *m) {
               } else {
                 return "arg_" + std::to_string(block_arg.index());
               }
+            } else if (self.first_use()) {
+              auto nextOp = self.first_use().owner();
+              if (nextOp->isa<::pir::ShadowOutputOp>()) {
+                return nextOp->attribute<pir::StrAttribute>("output_name")
+                    .AsString();
+              } else {
+                PADDLE_THROW(phi::errors::InvalidArgument(
+                    "Currently, we can only get name of Value which is "
+                    "shadowoutput "));
+              }
             } else {
               PADDLE_THROW(phi::errors::InvalidArgument(
                   "Currently, we can only get name of Value that "
                   "is persistable"));
             }
           })
+      .def_property_readonly(
+          "has_name",
+          [](Value self) {
+            if (self.defining_op()->isa<::pir::ParameterOp>() ||
+                self.defining_op()->isa<paddle::dialect::DataOp>() ||
+                self.isa<BlockArgument>() ||
+                (self.first_use() &&
+                 self.first_use().owner()->isa<::pir::ShadowOutputOp>())) {
+              return true;
+            } else {
+              return false;
+            }
+          })
       .def_property(
           "shape",
           [](Value self) { return phi::vectorize(GetValueDims(self)); },
 
@@ -60,6 +60,7 @@
     SymbolicStaticFunction,
     unwrap_decorators,
 )
+from .pir_translated_layer import PIR_INFER_MODEL_SUFFIX, PirTranslatedLayer
 from .translated_layer import (
     INFER_MODEL_SUFFIX,
     INFER_PARAMS_INFO_SUFFIX,
@@ -575,7 +576,11 @@ def _get_output_vars(outputs, output_spec, with_hook=False):
 def _build_load_path_and_config(path, config):
     # NOTE(chenweihang): If both [prefix save format] and [directory save format] exist,
     # raise error, avoid confusing behavior
-    prefix_format_path = path + INFER_MODEL_SUFFIX
+    if use_pir_api():
+        model_suffix = PIR_INFER_MODEL_SUFFIX
+    else:
+        model_suffix = INFER_MODEL_SUFFIX
+    prefix_format_path = path + model_suffix
     prefix_format_exist = os.path.exists(prefix_format_path)
     directory_format_exist = os.path.isdir(path)
     if prefix_format_exist and directory_format_exist:
@@ -600,7 +605,7 @@ def _build_load_path_and_config(path, config):
                     "specified file prefix, the ``model_filename`` config does "
                     "not take effect."
                 )
-            config.model_filename = file_prefix + INFER_MODEL_SUFFIX
+            config.model_filename = file_prefix + model_suffix
             if config.params_filename is not None:
                 warnings.warn(
                     "When loading the result saved with the "
@@ -1508,15 +1513,14 @@ def load(path, **configs):
                 ...         print("Epoch {} batch {}: loss = {}".format(
                 ...             epoch_id, batch_id, np.mean(loss.numpy())))
     """
-    if use_pir_api():
-        raise NotImplementedError(
-            "Currently, `paddle.jit.load` is not supported in PIR mode."
-        )
     # 1. construct correct config
     config = _parse_load_config(configs)
     model_path, config = _build_load_path_and_config(path, config)
 
-    return TranslatedLayer._construct(model_path, config)
+    if use_pir_api():
+        return PirTranslatedLayer._construct(model_path, config)
+    else:
+        return TranslatedLayer._construct(model_path, config)
 
 
 def set_dynamic_shape(variable, shape_list):