// If device attribute is not found from the function, nothing is done.
class TFDeviceAssignmentByFuncAttrPass
    : public impl::TFDeviceAssignmentByFuncAttrPassBase<
          TFDeviceAssignmentByFuncAttrPass> {
 public:
  TFDeviceAssignmentByFuncAttrPass() = default;
  TFDeviceAssignmentByFuncAttrPass(const TFDeviceAssignmentByFuncAttrPass&) {}

  void runOnOperation() override {
    func::FuncOp func = getOperation();

      %2 = "tf.Const"() {device = "baz", value = dense<[[42.0]]> : tensor<1x1xf32>} : () -> tensor<1x1xf32>
    ```
  }];
}

def TFDeviceAssignmentByFuncAttrPass : Pass<"tf-device-assignment-by-func-attr", "mlir::func::FuncOp"> {
  let summary = "Device assignment in TF dialect using the device specified in the function attribute.";