Fix vectorized MLIR derivative constants

enzyme/Enzyme/MLIR/Implementations/MathDerivatives.td

@@ -22,7 +22,7 @@ def : MLIRDerivative<"math", "SqrtOp", (Op $x),
                   >;
 def : MLIRDerivative<"math", "AtanOp", (Op $x),
                     [
-                      (CheckedMulF (DiffeRet), (DivF (ConstantFP<"1.0","arith","ConstantOp"> $x):$one, (AddF (MulF $x, $x), $one)))
+                      (CheckedMulF (DiffeRet), (DivF (ConstantFP<"1.0","arith","ConstantOp">):$one, (AddF (MulF $x, $x), $one)))
                     ]
                   >;
 def : MLIRDerivative<"math", "AbsFOp", (Op $x),

enzyme/test/MLIR/ForwardMode/batched_derivative_constant.mlir

@@ -0,0 +1,20 @@
+// RUN: %eopt --enzyme %s | FileCheck %s
+
+// Regression for a width-forward derivative with an operand-less floating
+// constant. The constant participates in the tangent expression and must use
+// the width-stacked shadow type.
+module {
+  func.func @atan(%x: tensor<4xf32>) -> tensor<4xf32> {
+    %0 = math.atan %x : tensor<4xf32>
+    return %0 : tensor<4xf32>
+  }
+
+  func.func @main(%x: tensor<4xf32>, %dx: tensor<2x4xf32>) -> tensor<2x4xf32> {
+    %0 = enzyme.fwddiff @atan(%x, %dx) {activity = [#enzyme<activity enzyme_dup>], ret_activity = [#enzyme<activity enzyme_dupnoneed>], width = 2 : i64} : (tensor<4xf32>, tensor<2x4xf32>) -> tensor<2x4xf32>
+    return %0 : tensor<2x4xf32>
+  }
+}
+
+// CHECK-LABEL: func.func private @fwddiffe2atan
+// CHECK: arith.constant dense<1.000000e+00> : tensor<2x4xf32>
+// CHECK: arith.divf

enzyme/tools/enzyme-tblgen/enzyme-tblgen.cpp

@@ -608,6 +608,7 @@ bool handle(const Twine &curIndent, const Twine &argPattern, raw_ostream &os,
         PrintFatalError(pattern->getLoc(), Twine("'value' not defined in ") +
                                                resultTree->getAsString());
 
+      bool vectorizedConstant = false;
       if (intrinsic == MLIRDerivatives) {
         if (resultRoot->getNumArgs() > 1)
           PrintFatalError(pattern->getLoc(),
@@ -618,7 +619,12 @@ bool handle(const Twine &curIndent, const Twine &argPattern, raw_ostream &os,
         std::string ord;
         bool shadowType = false;
         if (resultRoot->getNumArgs() == 0) {
-          ord = "op->getResult(0)";
+          ord = "gutils->getShadowType(op->getResult(0)";
+          // This constant participates in a tangent expression.  In vector
+          // forward mode, the tangent has a leading width dimension even when
+          // the primal result does not.  Use the shadow type to preserve the
+          // tangent shape.
+          shadowType = true;
         } else {
           if (resultRoot->getArgName(0)) {
             auto name = resultRoot->getArgName(0)->getAsUnquotedString();
@@ -665,6 +671,7 @@ bool handle(const Twine &curIndent, const Twine &argPattern, raw_ostream &os,
           os << ")";
         os << ", ";
         os << "\"" << value->getValue() << "\"))";
+        vectorizedConstant = shadowType;
       } else {
         if (resultRoot->getNumArgs() != 1)
           PrintFatalError(pattern->getLoc(),
@@ -695,7 +702,7 @@ bool handle(const Twine &curIndent, const Twine &argPattern, raw_ostream &os,
                               resultTree->getAsString());
         os << ", \"" << value->getValue() << "\")";
       }
-      return false;
+      return vectorizedConstant;
     } else if (opName == "Zero" || Def->isSubClassOf("Zero")) {
       if (resultRoot->getNumArgs() != 1)
         PrintFatalError(pattern->getLoc(), "only single op Zero supported");