Implement constant checks for symmetry

Author: gaurav-arya

src/enzyme_ad/jax/Utils.cpp

@@ -593,11 +593,45 @@ SymmetricResultAnalysis initSymmetricResultAnalysis() {
 }
 
 bool SymmetricResultAnalysis::constantIntCheck(DenseElementsAttr attr) {
-  return false; // TODO
+  auto type = dyn_cast<RankedTensorType>(attr.getType());
+  if (!type || type.getRank() != 2)
+    return false;
+
+  int64_t n = type.getDimSize(0);
+  int64_t m = type.getDimSize(1);
+  if (n != m)
+    return false;
+
+  auto values = attr.getValues<APInt>();
+  for (int64_t i = 0; i < n; ++i) {
+    for (int64_t j = 0; j < n; ++j) {
+      if (values[i * n + j] != values[j * n + i])
+        return false;
+    }
+  }
+
+  return true;
 }
 
 bool SymmetricResultAnalysis::constantFloatCheck(DenseElementsAttr attr) {
-  return false; // TODO
+  auto type = dyn_cast<RankedTensorType>(attr.getType());
+  if (!type || type.getRank() != 2)
+    return false;
+
+  int64_t n = type.getDimSize(0);
+  int64_t m = type.getDimSize(1);
+  if (n != m)
+    return false;
+
+  auto values = attr.getValues<APFloat>();
+  for (int64_t i = 0; i < n; ++i) {
+    for (int64_t j = 0; j < n; ++j) {
+      if (values[i * n + j] != values[j * n + i])
+        return false;
+    }
+  }
+
+  return true;
 }
 
 SymmetricResultAnalysis::State SymmetricResultAnalysis::localGuaranteed(