diff --git a/.github/workflows/Benchmark.yml b/.github/workflows/Benchmark.yml
index 61fa42efd..5bb1f92ee 100644
--- a/.github/workflows/Benchmark.yml
+++ b/.github/workflows/Benchmark.yml
@@ -2,7 +2,7 @@ name: Benchmarks
 on:
   pull_request:
     branches:
-      - master
+      - main
 
 concurrency:
   # Skip intermediate builds: always.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index dcd006c0b..06219ec9f 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -2,7 +2,7 @@ name: CI
 on:
   push:
     branches:
-      - master
+      - main
     tags: ['*']
   pull_request:
   workflow_dispatch:
diff --git a/.github/workflows/Docs.yml b/.github/workflows/Docs.yml
index 35de66204..9128876b5 100644
--- a/.github/workflows/Docs.yml
+++ b/.github/workflows/Docs.yml
@@ -3,7 +3,7 @@ name: Documentation
 on:
   push:
     branches:
-      - master
+      - main
     tags: '*'
   pull_request:
 
diff --git a/.github/workflows/Enzyme.yml b/.github/workflows/Enzyme.yml
index d3dbf1fc9..de16a19d0 100644
--- a/.github/workflows/Enzyme.yml
+++ b/.github/workflows/Enzyme.yml
@@ -2,7 +2,7 @@ name: Enzyme
 on:
   push:
     branches:
-      - master
+      - main
     tags: ['*']
   pull_request:
   workflow_dispatch:
diff --git a/Project.toml b/Project.toml
index 5f44dadda..3dede7598 100644
--- a/Project.toml
+++ b/Project.toml
@@ -23,10 +23,6 @@ StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 [weakdeps]
 Bijectors = "76274a88-744f-5084-9051-94815aaf08c4"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
-ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
-Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
-ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
 AdvancedVIBijectorsExt = "Bijectors"
@@ -43,18 +39,14 @@ Distributions = "0.25.111"
 DocStringExtensions = "0.8, 0.9"
 Enzyme = "0.13"
 FillArrays = "1.3"
-ForwardDiff = "0.10"
 Functors = "0.4, 0.5"
 LinearAlgebra = "1"
 LogDensityProblems = "2"
-Mooncake = "0.4"
 Optimisers = "0.2.16, 0.3, 0.4"
 ProgressMeter = "1.6"
 Random = "1"
 Requires = "1.0"
-ReverseDiff = "1"
 StatsBase = "0.32, 0.33, 0.34"
-Zygote = "0.6"
 julia = "1.10, 1.11.2"
 
 [extras]
diff --git a/README.md b/README.md
index e5d1ea7f2..8949cf4ee 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,7 @@
 [![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://turinglang.org/AdvancedVI.jl/stable/)
 [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://turinglang.org/AdvancedVI.jl/dev/)
-[![Build Status](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/CI.yml/badge.svg?branch=master)](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/CI.yml?query=branch%3Amaster)
-[![Coverage](https://codecov.io/gh/TuringLang/AdvancedVI.jl/branch/master/graph/badge.svg)](https://codecov.io/gh/TuringLang/AdvancedVI.jl)
+[![Build Status](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/TuringLang/AdvancedVI.jl/actions/workflows/CI.yml?query=branch%3Amain)
+[![Coverage](https://codecov.io/gh/TuringLang/AdvancedVI.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/TuringLang/AdvancedVI.jl)
 
 # AdvancedVI.jl
 
diff --git a/src/optimize.jl b/src/optimize.jl
index 5c65ad777..69be56390 100644
--- a/src/optimize.jl
+++ b/src/optimize.jl
@@ -58,7 +58,7 @@ function optimize(
     adtype::ADTypes.AbstractADType,
     optimizer::Optimisers.AbstractRule=Optimisers.Adam(),
     averager::AbstractAverager=NoAveraging(),
-    operator::AbstractOperator=IdentityOperator(),
+    operator=IdentityOperator(),
     show_progress::Bool=true,
     state_init::NamedTuple=NamedTuple(),
     callback=nothing,