test: exported functions

Author: avik-pal

CondaPkg.toml

@@ -1,7 +1,7 @@
 [deps]
-python = "<=3.13,>=3.9,<4"
+python = "<=3.12,>=3.9,<4"
 
 [pip.deps]
-jax = ">= 0.6"
+jax = ">= 0.5"
 tensorflow = ">= 2.17"
-numpy = ">= 2"
+numpy = ">= 1, >= 2"

Project.toml

@@ -98,6 +98,7 @@ Libdl = "1.10"
 LinearAlgebra = "1.10"
 MPI = "0.20"
 NNlib = "0.9.26"
+NPZ = "0.4"
 OffsetArrays = "1"
 OneHotArrays = "0.2.10"
 OrderedCollections = "1"

src/serialization/EnzymeJAX.jl

@@ -1,6 +1,6 @@
 module EnzymeJAX
 
-using ..Reactant: Reactant, Compiler, MLIR, Serialization
+using ..Reactant: Reactant, Compiler, Serialization
 
 """
     export_to_enzymejax(

test/Project.toml

@@ -3,6 +3,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 DLFP8Types = "f4c16678-4a16-415b-82ef-ed337c5d6c7c"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
@@ -18,9 +19,10 @@ KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
-MethodAnalysis = "85b6ec6f-f7df-4429-9514-a64bcd9ee824"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
+MethodAnalysis = "85b6ec6f-f7df-4429-9514-a64bcd9ee824"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
+NPZ = "15e1cf62-19b3-5cfa-8e77-841668bca605"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 OneHotArrays = "0b1bfda6-eb8a-41d2-88d8-f5af5cad476f"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
@@ -56,6 +58,7 @@ Lux = "1.21"
 LuxLib = "1.11"
 MPI = "0.20"
 NNlib = "0.9.26"
+NPZ = "0.4"
 OffsetArrays = "1"
 OneHotArrays = "0.2.6"
 Optimisers = "0.4"