From c5d326666dfa714ee8d56cac172c13838c4d7f44 Mon Sep 17 00:00:00 2001
From: Avik Pal <avikpal@mit.edu>
Date: Mon, 26 Jun 2023 11:45:31 -0400
Subject: [PATCH] Use PackageExtensionCompat

---
 .buildkite/pipeline.yml      |  5 -----
 .github/workflows/CI.yml     |  2 --
 Project.toml                 |  8 ++++----
 ext/BoltzFluxMetalheadExt.jl |  3 +--
 src/Boltz.jl                 | 10 ++--------
 5 files changed, 7 insertions(+), 21 deletions(-)

diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
index 09281ca..952aea8 100644
--- a/.buildkite/pipeline.yml
+++ b/.buildkite/pipeline.yml
@@ -21,13 +21,8 @@ steps:
       setup:
         julia:
           - "1"
-          - "1.6"
-          - "1.9-nightly"
           - "nightly"
       adjustments:
-        - with:
-            julia: "1.6"  # GPU compilation for some of the gigantic models break
-          soft_fail: true
         - with:
             julia: "nightly"
           soft_fail: true
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 697a2bd..cab3a0e 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -20,7 +20,6 @@ jobs:
         version:
           - "1"
           - "1.6"
-          - "~1.9.0-0"
     steps:
       - uses: actions/checkout@v3
       - uses: julia-actions/setup-julia@v1
@@ -44,4 +43,3 @@ jobs:
       - uses: codecov/codecov-action@v3
         with:
           files: lcov.info
-          flags: ${{ matrix.group }}
diff --git a/Project.toml b/Project.toml
index 966fbb9..fb226dd 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "Boltz"
 uuid = "4544d5e4-abc5-4dea-817f-29e4c205d9c8"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 Artifacts = "56f22d72-fd6d-98f1-02f0-08ddc0907c33"
@@ -11,8 +11,8 @@ JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LazyArtifacts = "4af54fe1-eca0-43a8-85a7-787d91b784e3"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+PackageExtensionCompat = "65ce6f38-6b18-4e1d-a461-8949797d7930"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [weakdeps]
@@ -27,8 +27,8 @@ ChainRulesCore = "1.15"
 JLD2 = "0.4"
 Lux = "0.4.26"
 Metalhead = "0.7"
-NNlib = "0.8"
-Requires = "1"
+NNlib = "0.8, 0.9"
+PackageExtensionCompat = "1"
 julia = "1.6"
 
 [extras]
diff --git a/ext/BoltzFluxMetalheadExt.jl b/ext/BoltzFluxMetalheadExt.jl
index 6ff1bb2..4a4993a 100644
--- a/ext/BoltzFluxMetalheadExt.jl
+++ b/ext/BoltzFluxMetalheadExt.jl
@@ -1,7 +1,6 @@
 module BoltzFluxMetalheadExt
 
-isdefined(Base, :get_extension) ? (using Metalhead) : (using ..Metalhead)
-using Boltz, Lux
+using Boltz, Lux, Metalhead
 import Boltz: alexnet, convmixer, densenet, googlenet, mobilenet, resnet, resnext
 using Boltz: _initialize_model, assert_name_present_in
 
diff --git a/src/Boltz.jl b/src/Boltz.jl
index 0a43255..c1b854a 100644
--- a/src/Boltz.jl
+++ b/src/Boltz.jl
@@ -7,15 +7,9 @@ using Artifacts, JLD2, LazyArtifacts
 import ChainRulesCore as CRC
 
 # Extensions
-if !isdefined(Base, :get_extension)
-    using Requires
-end
-
+using PackageExtensionCompat
 function __init__()
-    @static if !isdefined(Base, :get_extension)
-        # Import Flux Models
-        @require Metalhead="dbeba491-748d-5e0e-a39e-b530a07fa0cc" begin include("../ext/BoltzFluxMetalheadExt.jl") end
-    end
+    @require_extensions
 end
 
 # Define functions. Methods defined in files or in extensions later