Index: mlir/lib/AsmParser/AttributeParser.cpp
===================================================================
--- mlir/lib/AsmParser/AttributeParser.cpp
+++ mlir/lib/AsmParser/AttributeParser.cpp
@@ -1241,13 +1241,20 @@
   }
 
   // Parse the referenced attribute.
-  if (parseToken(Token::r_square, "expected ']' to close distinct ID") ||
-      parseToken(Token::less, "expected '<' after distinct ID"))
-    return {};
-  Attribute referencedAttr = parseAttribute(type);
-  if (!referencedAttr) {
-    emitError("expected attribute");
+  if (parseToken(Token::r_square, "expected ']' to close distinct ID"))
     return {};
+
+  Attribute referencedAttr = builder.getUnitAttr();
+  if (getToken().is(Token::less)) {
+    consumeToken();
+    referencedAttr = parseAttribute(type);
+    if (!referencedAttr) {
+      emitError("expected attribute");
+      return {};
+    }
+
+    if (parseToken(Token::greater, "expected '>' to close distinct attribute"))
+      return {};
   }
 
   // Add the distinct attribute to the parser state, if it has not been parsed
@@ -1265,8 +1272,5 @@
     return {};
   }
 
-  if (parseToken(Token::greater, "expected '>' to close distinct attribute"))
-    return {};
-
   return it->getSecond();
 }
Index: mlir/lib/IR/AsmPrinter.cpp
===================================================================
--- mlir/lib/IR/AsmPrinter.cpp
+++ mlir/lib/IR/AsmPrinter.cpp
@@ -2140,9 +2140,15 @@
     os << "unit";
     return;
   } else if (auto distinctAttr = llvm::dyn_cast<DistinctAttr>(attr)) {
-    os << "distinct[" << state.getDistinctState().getId(distinctAttr) << "]<";
+    os << "distinct[" << state.getDistinctState().getId(distinctAttr) << "]";
+    if (llvm::isa<UnitAttr>(distinctAttr.getReferencedAttr())) {
+      return;
+    }
+
+    os << '<';
     printAttribute(distinctAttr.getReferencedAttr());
     os << '>';
+
     return;
   } else if (auto dictAttr = llvm::dyn_cast<DictionaryAttr>(attr)) {
     os << '{';
Index: mlir/lib/IR/BuiltinDialect.cpp
===================================================================
--- mlir/lib/IR/BuiltinDialect.cpp
+++ mlir/lib/IR/BuiltinDialect.cpp
@@ -60,10 +60,11 @@
       os << "loc";
       return AliasResult::OverridableAlias;
     }
-    if (llvm::isa<DistinctAttr>(attr)) {
-      os << "distinct";
-      return AliasResult::OverridableAlias;
-    }
+    if (auto distinct = llvm::dyn_cast<DistinctAttr>(attr))
+      if (!llvm::isa<UnitAttr>(distinct.getReferencedAttr())) {
+        os << "distinct";
+        return AliasResult::OverridableAlias;
+      }
     return AliasResult::NoAlias;
   }
 
Index: mlir/test/IR/distinct-attr.mlir
===================================================================
--- mlir/test/IR/distinct-attr.mlir
+++ mlir/test/IR/distinct-attr.mlir
@@ -20,3 +20,10 @@
 // CHECK:         distinct.attr = #[[DISTINCT2]]
 // CHECK-GENERIC: distinct.attr = distinct[2]<42 : i32>
 "test.op"() {distinct.attr = distinct[42]<42 : i32>} : () -> ()
+
+// CHECK: distinct.attr = distinct[3]
+// CHECK-NOT: <unit>
+"test.op"() {distinct.attr = distinct[3]} : () -> ()
+
+// CHECK: distinct.attr = distinct[4]
+"test.op"() {distinct.attr = distinct[4]<unit>} : () -> ()
Index: mlir/test/IR/invalid-builtin-attributes.mlir
===================================================================
--- mlir/test/IR/invalid-builtin-attributes.mlir
+++ mlir/test/IR/invalid-builtin-attributes.mlir
@@ -569,11 +569,6 @@
 
 // -----
 
-// expected-error@below {{expected '<' after distinct ID}}
-#attr = distinct[8](
-
-// -----
-
 // expected-error@below {{expected attribute}}
 #attr = distinct[8]<attribute