Support retrieving the splat value from DenseElementsAttrs in Python

[apaszke]

Summary

This is especially convenient when trying to resize the splat.

[jpienaar]

Looks good thanks

[rkayaith]

mlir/lib/Bindings/Python/IRAttributes.cpp

782–786

nit:

mlir/lib/Bindings/Python/IRAttributes.cpp

c.def("__len__", &PyDenseElementsAttribute::dunderLen)
                     kDenseElementsAttrGetDocstring)
         .def_static("get_splat", PyDenseElementsAttribute::getSplat,
                     py::arg("shaped_type"), py::arg("element_attr"),
                     "Gets a DenseElementsAttr where all values are the same")
         .def_property_readonly("is_splat",
                                [](PyDenseElementsAttribute &self) -> bool {
                                  return mlirDenseElementsAttrIsSplat(self);
                                })
+        .def("get_splat_value",
+             [](PyDenseElementsAttribute &self) -> PyAttribute {
+               if (!mlirDenseElementsAttrIsSplat(self)) {
+                 throw SetPyError(
+                     PyExc_ValueError,
+                     "get_splat_value called on a non-splat attribute");
+               }
+               return PyAttribute(self.getContext(),
+                                  mlirDenseElementsAttrGetSplatValue(self));
+             })
         .def_buffer(&PyDenseElementsAttribute::accessBuffer);
   }
 
 private:
   static bool isUnsignedIntegerFormat(const std::string &format) {
     if (format.empty())
       return false;
     char code = format[0];

mlir/test/python/ir/array_attributes.py

with Context(), Location.unknown():
     t = IntegerType.get_signless(32)
     element = IntegerAttr.get(t, 555)
     shaped_type = RankedTensorType.get((2, 3, 4), t)
     attr = DenseElementsAttr.get_splat(shaped_type, element)
     # CHECK: dense<555> : tensor<2x3x4xi32>
     print(attr)
     # CHECK: is_splat: True
     print("is_splat:", attr.is_splat)
+    assert attr.get_splat_value() == element
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsSplatFloat
 @run
 def testGetDenseElementsSplatFloat():
   with Context(), Location.unknown():
     t = F32Type.get()
     element = FloatAttr.get(t, 1.2)
     shaped_type = RankedTensorType.get((2, 3, 4), t)
     attr = DenseElementsAttr.get_splat(shaped_type, element)
     # CHECK: dense<1.200000e+00> : tensor<2x3x4xf32>
     print(attr)
+    assert attr.get_splat_value() == element
 
 
 # CHECK-LABEL: TEST: testGetDenseElementsSplatErrors
 @run
 def testGetDenseElementsSplatErrors():
   with Context(), Location.unknown():
     t = F32Type.get()
     other_t = F64Type.get()