Deprecate ShapeOfXlaOp in favor of GetShape (#9381)

Co-authored-by: Zhanyong Wan <[email protected]>

Author: zhanyong-wan

torch_xla/csrc/BUILD

@@ -243,6 +243,8 @@ ptxla_cc_library(
     srcs = ["shape_builder.cpp"],
     hdrs = ["shape_builder.h"],
     deps = [
+        "@com_google_absl//absl/base:core_headers",
+        "@com_google_absl//absl/base:nullability",
         "@com_google_absl//absl/types:span",
         "@xla//xla:shape_util",
         "@xla//xla:types",

torch_xla/csrc/shape_helper.cpp

@@ -6,8 +6,11 @@
 namespace torch_xla {
 
 const xla::Shape& ShapeHelper::ShapeOfXlaOp(xla::XlaOp op) {
-  const xla::Shape* shape = ConsumeValue(op.builder()->GetShapePtr(op));
-  return *shape;
+  return *ConsumeValue(GetShape(op));
+}
+
+absl::StatusOr<const xla::Shape * absl_nonnull> GetShape(xla::XlaOp op) {
+  return op.builder()->GetShapePtr(op);
 }
 
 }  // namespace torch_xla

torch_xla/csrc/shape_helper.h

@@ -1,16 +1,24 @@
 #ifndef XLA_TORCH_XLA_SHAPE_HELPER_H_
 #define XLA_TORCH_XLA_SHAPE_HELPER_H_
 
+#include "absl/base/attributes.h"
+#include "absl/base/nullability.h"
 #include "xla/hlo/builder/xla_builder.h"
 
 namespace torch_xla {
 
 class ShapeHelper {
  public:
   // Returns the shape of the given XLA operation.
+  ABSL_DEPRECATED(
+      "Use GetShape(op) instead. ShapeOfXlaOp() blindly "
+      "de-references StatusOr returned by XLA, which is unsafe.")
   static const xla::Shape& ShapeOfXlaOp(xla::XlaOp op);
 };
 
+// Returns the shape of the given XLA operation.
+absl::StatusOr<const xla::Shape * absl_nonnull> GetShape(xla::XlaOp op);
+
 }  // namespace torch_xla
 
 #endif  // XLA_TORCH_XLA_SHAPE_HELPER_H_