Bump xla and fix ci

Author: wsmoses

.buildkite/pipeline.yml

@@ -7,6 +7,8 @@ steps:
           - x86_64
         os:
           - macos
+        python:
+          - "3.10"
     agents:
       queue: "juliaecosystem"
       os: "{{matrix.os}}"
@@ -71,7 +73,7 @@ steps:
       python -m pip install --user numpy wheel
       mkdir -p .baztmp
       rm -f bazel-bin/*.whl
-      bazel --output_user_root=`pwd`/.baztmp build :enzyme_ad
+      HERMETIC_PYTHON_VERSION=${{matrix.python}} bazel --output_user_root=`pwd`/.baztmp build :enzyme_ad
       cp bazel-bin/*.whl .
       python -m pip install *.whl
       python -m pip install "jax[cpu]"

.buildkite/secure_pipeline.yml

@@ -7,6 +7,11 @@ steps:
           - x86_64
         os:
           - macos
+        python:
+          - "3.10"
+          - "3.11"
+          - "3.12"
+          - "3.13"
     agents:
       queue: "juliaecosystem"
       os: "{{matrix.os}}"
@@ -59,18 +64,15 @@ steps:
           wget https://github.com/bazelbuild/bazel/releases/download/6.2.1/bazel-6.2.1-windows-x86_64.exe
         fi
         mv bazel* .local/bin/bazel.exe
-        start /wait "" Miniconda3*.exe /InstallationType=JustMe /RegisterPython=0 /S /D=`pwd`/conda
-        rm Miniconda*.exe
       fi
       python -m ensurepip --upgrade
       python -m pip install --user numpy wheel
       mkdir baztmp
       export TAG=`echo $BUILDKITE_TAG | cut -c2-`
       sed -i.bak "s~version = \"[0-9.]*\"~version = \"\$TAG\"~g" BUILD
-      bazel --output_user_root=`pwd`/baztmp build :enzyme_ad
+      HERMETIC_PYTHON_VERSION=${{matrix.python}} bazel --output_user_root=`pwd`/baztmp build :enzyme_ad
       cp bazel-bin/*.whl .
       python -m pip install *.whl
-      cd test && python -m pip install "jax[cpu]" && python test.py && cd ..
       python -m pip install --user twine
       python -m twine upload *.whl
     artifact_paths:

src/enzyme_ad/jax/BUILD

@@ -258,7 +258,6 @@ cc_library(
         "@llvm-project//mlir:Pass",
         "@llvm-project//mlir:Support",
         "@llvm-project//mlir:TensorDialect",
-        "@llvm-project//mlir:Transforms",
         "@llvm-project//mlir:TransformUtils",
         "@stablehlo//:reference_ops",
         "@stablehlo//:stablehlo_ops",

src/enzyme_ad/jax/Implementations/HLODerivatives.td

@@ -143,7 +143,7 @@ def PadToSliceStride : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, [{
 
 def ResultDotDim : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, "op.getDotDimensionNumbersAttr()">;
 def ResultDotPrec : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, "op.getPrecisionConfigAttr()">;
-
+def ResultDotAlg : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, "op.getAlgorithmAttr()">;
 
 def ShadowLHSDotDim : GlobalExpr</*needsprimal*/0, /*needsshadow*/0, [{
   auto existingattr = op.getDotDimensionNumbersAttr();
@@ -572,10 +572,10 @@ def : HLODerivative<"DivOp", (Op $x, $y),
 
 def : HLODerivative<"DotGeneralOp", (Op $lhs, $rhs),
                     [
-                        (Transpose (TypeOf $lhs), (DotGeneral (ShadowLHSDotRes), (DiffeRet), $rhs, (ShadowLHSDotDim), (ResultDotPrec)), (ShadowLHSTranspose)),
-                        (Transpose (TypeOf $rhs), (DotGeneral (ShadowRHSDotRes), $lhs, (DiffeRet), (ShadowRHSDotDim), (ResultDotPrec)), (ShadowRHSTranspose))
+                        (Transpose (TypeOf $lhs), (DotGeneral (ShadowLHSDotRes), (DiffeRet), $rhs, (ShadowLHSDotDim), (ResultDotPrec), (ResultDotAlg)), (ShadowLHSTranspose)),
+                        (Transpose (TypeOf $rhs), (DotGeneral (ShadowRHSDotRes), $lhs, (DiffeRet), (ShadowRHSDotDim), (ResultDotPrec), (ResultDotAlg)), (ShadowRHSTranspose))
                       ],
-                    (Add (SelectIfActive $lhs, (DotGeneral (ResultTypes), (Shadow $lhs), $rhs, (ResultDotDim), (ResultDotPrec)), (HLOConstantFP<"0">)), (SelectIfActive $rhs, (DotGeneral (ResultTypes), $lhs, (Shadow $rhs), (ResultDotDim), (ResultDotPrec)), (HLOConstantFP<"0">)))
+                    (Add (SelectIfActive $lhs, (DotGeneral (ResultTypes), (Shadow $lhs), $rhs, (ResultDotDim), (ResultDotPrec), (ResultDotAlg)), (HLOConstantFP<"0">)), (SelectIfActive $rhs, (DotGeneral (ResultTypes), $lhs, (Shadow $rhs), (ResultDotDim), (ResultDotPrec), (ResultDotAlg)), (HLOConstantFP<"0">)))
                   >;
 
 def : HLOInactiveOp<"DynamicIotaOp">;

src/enzyme_ad/jax/Passes/EnzymeHLOOpt.cpp

@@ -3590,13 +3590,13 @@ struct TransposeDotReorder
     if (!convert) {
       rewriter.replaceOpWithNewOp<stablehlo::DotGeneralOp>(
           op, op.getType(), dot.getRhs(), dot.getLhs(), ndim,
-          dot.getPrecisionConfigAttr());
+          dot.getPrecisionConfigAttr(), dot.getAlgorithmAttr());
     } else {
       auto middot = rewriter.create<stablehlo::DotGeneralOp>(
           op.getLoc(),
           RankedTensorType::get(op.getType().getShape(),
                                 dot.getType().getElementType()),
-          dot.getRhs(), dot.getLhs(), ndim, dot.getPrecisionConfigAttr());
+          dot.getRhs(), dot.getLhs(), ndim, dot.getPrecisionConfigAttr(), dot.getAlgorithmAttr());
       rewriter.replaceOpWithNewOp<stablehlo::ConvertOp>(op, op.getType(),
                                                         middot);
     }
@@ -3808,7 +3808,7 @@ struct DotTranspose : public OpRewritePattern<mlir::stablehlo::DotGeneralOp> {
     rewriter.replaceOpWithNewOp<stablehlo::DotGeneralOp>(
         dot, dot.getType(), lhs_trans ? lhs_trans.getOperand() : dot.getLhs(),
         rhs_trans ? rhs_trans.getOperand() : dot.getRhs(), ndim,
-        dot.getPrecisionConfigAttr());
+        dot.getPrecisionConfigAttr(), dot.getAlgorithmAttr());
     return success();
   }
 };
@@ -4712,7 +4712,8 @@ struct PadDotGeneral : public OpRewritePattern<mlir::stablehlo::DotGeneralOp> {
         RankedTensorType::get(resultShape, op.getType().getElementType()),
         otherIsLHS ? nextOtherArg : pad.getOperand(),
         otherIsLHS ? pad.getOperand() : nextOtherArg,
-        op.getDotDimensionNumbersAttr(), op.getPrecisionConfigAttr());
+        op.getDotDimensionNumbersAttr(), op.getPrecisionConfigAttr(),
+        op.getAlgorithmAttr(), op.getAlgorithmAttr());
 
     if (!resultDimsToPad.empty()) {
       SmallVector<int64_t> low(op.getType().getShape().size(), 0);

src/enzyme_ad/jax/enzyme_call.cc

@@ -902,7 +902,7 @@ class CpuKernel {
                       pyargv, mode, lang, xla_runtime, pass_pipeline);
 
     if (!JIT) {
-      DL = std::make_unique<llvm::DataLayout>(mod.get());
+      DL = std::make_unique<llvm::DataLayout>(mod->getDataLayoutStr());
       auto tJIT =
           llvm::orc::LLJITBuilder()
               .setDataLayout(*DL.get())

workspace.bzl

@@ -1,7 +1,7 @@
-JAX_COMMIT = "734ebd570891ceaf8c7104e12256a1edfe942b14"
-JAX_SHA256 = "ef7b578520f7fdc189ac2579f2929c81d2a6bd07b7a1e1cfa0017111fc0d12ca"
+JAX_COMMIT = "810a91968a853b4ae15aa5c5282e5673136bb980"
+JAX_SHA256 = ""
 
-ENZYME_COMMIT = "6e8dd3e3faff7e766a9e957bf4068d6fcda54539"
+ENZYME_COMMIT = "cc65abdb55e5e2e142d773e0508e1083d4b3ac52"
 ENZYME_SHA256 = ""
 
 XLA_PATCHES = [