almost done here

Project.toml

@@ -14,27 +14,35 @@ MLJModelInterface = "e80e1ace-859a-464e-9ed9-23947d8ae3ea"
 MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54"
 ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Tables = "bd369af6-aec1-5ad0-b16a-f7cc5008161c"
+TaijaBase = "10284c91-9f28-4c9a-abbf-ee43576dfff6"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
+Aqua = "0.8"
 CategoricalArrays = "0.10"
 ChainRulesCore = "1.16"
 ComputationalResources = "0.3"
 Distributions = "0.25"
 Flux = "0.13, 0.14"
-MLJFlux = "0.2, 0.3"
+MLJFlux = "0.2, 0.3, 0.4.0"
 MLJModelInterface = "1.8"
 MLUtils = "0.4"
 ProgressMeter = "1.7"
+Reexport = "1.2.2"
 StatsBase = "0.33, 0.34"
 Tables = "1.10"
+TaijaBase = "1.1.0"
 Zygote = "0.6"
-julia = "1.7"
+Random = "1.7, 1.10"
+Test = "1.7, 1.10"
+julia = "1.7, 1.10"
 
 [extras]
+Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Aqua", "Test"]

README.md

@@ -3,7 +3,7 @@
 
 *Joint Energy Models in Julia.*
 
-[![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://juliatrustworthyai.github.io/JointEnergyModels.jl/stable) [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://juliatrustworthyai.github.io/JointEnergyModels.jl/dev) [![Build Status](https://github.com/juliatrustworthyai/JointEnergyModels.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/juliatrustworthyai/JointEnergyModels.jl/actions/workflows/CI.yml?query=branch%3Amain) [![Coverage](https://codecov.io/gh/juliatrustworthyai/JointEnergyModels.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/juliatrustworthyai/JointEnergyModels.jl) [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle) [![License](https://img.shields.io/github/license/juliatrustworthyai/JointEnergyModels.jl)](LICENSE) [![Package Downloads](https://shields.io/endpoint?url=https://pkgs.genieframework.com/api/v1/badge/JointEnergyModels/.png)](https://pkgs.genieframework.com?packages=JointEnergyModels)
+[![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://juliatrustworthyai.github.io/JointEnergyModels.jl/stable) [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://juliatrustworthyai.github.io/JointEnergyModels.jl/dev) [![Build Status](https://github.com/juliatrustworthyai/JointEnergyModels.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/juliatrustworthyai/JointEnergyModels.jl/actions/workflows/CI.yml?query=branch%3Amain) [![Coverage](https://codecov.io/gh/juliatrustworthyai/JointEnergyModels.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/juliatrustworthyai/JointEnergyModels.jl) [![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle) [![License](https://img.shields.io/github/license/juliatrustworthyai/JointEnergyModels.jl)](LICENSE) [![Package Downloads](https://img.shields.io/badge/dynamic/json?url=http%3A%2F%2Fjuliapkgstats.com%2Fapi%2Fv1%2Fmonthly_downloads%2FJointEnergyModels&query=total_requests&suffix=%2Fmonth&label=Downloads)](http://juliapkgstats.com/pkg/JointEnergyModels) [![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)
 
 `JointEnergyModels.jl` is a package for training Joint Energy Models in Julia. Joint Energy Models (JEM) are hybrid models that learn to discriminate between classes $y$ and generate input data $x$. They were introduced in Grathwohl et al. (2020), which provides the foundation for the methodologies implemented in this package.
 

_freeze/docs/src/index/execute-results/md.json

@@ -2,7 +2,7 @@
   "hash": "8d26f827b22cadaf45d98344b37a0051",
   "result": {
     "engine": "jupyter",
-    "markdown": "```@meta\nCurrentModule = JointEnergyModels\n```\n\n# `JointEnergyModels.jl`\n\nDocumentation for [JointEnergyModels.jl](https://github.com/juliatrustworthyai/JointEnergyModels.jl).\n\n---\nexecute: \n  eval: false\n---\n\n*Joint Energy Models in Julia.*\n\n[![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://juliatrustworthyai.github.io/JointEnergyModels.jl/stable)\n[![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://juliatrustworthyai.github.io/JointEnergyModels.jl/dev)\n[![Build Status](https://github.com/juliatrustworthyai/JointEnergyModels.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/juliatrustworthyai/JointEnergyModels.jl/actions/workflows/CI.yml?query=branch%3Amain)\n[![Coverage](https://codecov.io/gh/juliatrustworthyai/JointEnergyModels.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/juliatrustworthyai/JointEnergyModels.jl)\n[![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle)\n[![License](https://img.shields.io/github/license/juliatrustworthyai/JointEnergyModels.jl)](LICENSE)\n[![Package Downloads](https://shields.io/endpoint?url=https://pkgs.genieframework.com/api/v1/badge/JointEnergyModels/)](https://pkgs.genieframework.com?packages=JointEnergyModels)\n\n\n\n`JointEnergyModels.jl` is a package for training Joint Energy Models in Julia. Joint Energy Models (JEM) are hybrid models that learn to discriminate between classes $y$ and generate input data $x$. They were introduced in @grathwohl2020your, which provides the foundation for the methodologies implemented in this package.\n\n## 🔁 Status\n\nThis package is still in its infancy and the API is subject to change. Currently, the package can be used to train JEMs for classification. It is also possible to train pure Energy-Based Models (EBMs) for the generative task only. The package is compatible with `Flux.jl`. Work on compatibility with `MLJ.jl` (through `MLJFlux.jl`) is currently under way.\n\nWe welcome contributions and feedback at this early stage. To install the development version of the package you can run the following command:\n\n```{.julia}\nusing Pkg\nPkg.add(url=\"https://github.com/juliatrustworthyai/JointEnergyModels.jl\")\n```\n\n## 🔍 Usage Example\n\nBelow we first generate some synthetic data:\n\n::: {.cell execution_count=2}\n``` {.julia .cell-code}\nnobs=2000\nX, y = make_circles(nobs, noise=0.1, factor=0.5)\nXplot = Float32.(permutedims(matrix(X)))\nX = table(permutedims(Xplot))\nplt = scatter(Xplot[1,:], Xplot[2,:], group=y, label=\"\")\nbatch_size = Int(round(nobs/10))\ndisplay(plt)\n```\n\n::: {.cell-output .cell-output-display}\n![](index_files/figure-commonmark/cell-3-output-1.svg){}\n:::\n:::\n\n\nThe `MLJ` compatible classifier can be instantiated as follows:\n\n::: {.cell execution_count=3}\n``` {.julia .cell-code}\n𝒟x = Normal()\n𝒟y = Categorical(ones(2) ./ 2)\nsampler = ConditionalSampler(𝒟x, 𝒟y, input_size=size(Xplot)[1:end-1], batch_size=batch_size)\nclf = JointEnergyClassifier(\n    sampler;\n    builder=MLJFlux.MLP(hidden=(32, 32, 32,), σ=Flux.relu),\n    batch_size=batch_size,\n    finaliser=x -> x,\n    loss=Flux.Losses.logitcrossentropy,\n)\n```\n:::\n\n\nIt uses the `MLJFlux` package to build the model:\n\n::: {.cell execution_count=4}\n``` {.julia .cell-code}\nprintln(typeof(clf) <: MLJFlux.MLJFluxModel)\n```\n\n::: {.cell-output .cell-output-stdout}\n```\ntrue\n```\n:::\n:::\n\n\nThe model can be wrapped in data and trained using the `fit!` function:\n\n::: {.cell execution_count=5}\n``` {.julia .cell-code}\nmach = machine(clf, X, y)\nfit!(mach)\n```\n:::\n\n\nThe results are visualised below. The model has learned to discriminate between the two classes (as indicated by the contours) and to generate samples from each class (as indicated by the stars).\n\n\n\n::: {.cell execution_count=7}\n\n::: {.cell-output .cell-output-display}\n![](index_files/figure-commonmark/cell-8-output-1.svg){}\n:::\n:::\n\n\n## 🎓 References\n\n",
+    "markdown": "```@meta\nCurrentModule = JointEnergyModels\n```\n\n# `JointEnergyModels.jl`\n\nDocumentation for [JointEnergyModels.jl](https://github.com/juliatrustworthyai/JointEnergyModels.jl).\n\n---\nexecute: \n  eval: false\n---\n\n*Joint Energy Models in Julia.*\n\n[![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://juliatrustworthyai.github.io/JointEnergyModels.jl/stable)\n[![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://juliatrustworthyai.github.io/JointEnergyModels.jl/dev)\n[![Build Status](https://github.com/juliatrustworthyai/JointEnergyModels.jl/actions/workflows/CI.yml/badge.svg?branch=main)](https://github.com/juliatrustworthyai/JointEnergyModels.jl/actions/workflows/CI.yml?query=branch%3Amain)\n[![Coverage](https://codecov.io/gh/juliatrustworthyai/JointEnergyModels.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/juliatrustworthyai/JointEnergyModels.jl)\n[![Code Style: Blue](https://img.shields.io/badge/code%20style-blue-4495d1.svg)](https://github.com/invenia/BlueStyle)\n[![License](https://img.shields.io/github/license/juliatrustworthyai/JointEnergyModels.jl)](LICENSE)\n[![Package Downloads](https://img.shields.io/badge/dynamic/json?url=http%3A%2F%2Fjuliapkgstats.com%2Fapi%2Fv1%2Fmonthly_downloads%2FJointEnergyModels&query=total_requests&suffix=%2Fmonth&label=Downloads)](http://juliapkgstats.com/pkg/JointEnergyModels) \n[![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)\n\n\n\n`JointEnergyModels.jl` is a package for training Joint Energy Models in Julia. Joint Energy Models (JEM) are hybrid models that learn to discriminate between classes $y$ and generate input data $x$. They were introduced in @grathwohl2020your, which provides the foundation for the methodologies implemented in this package.\n\n## 🔁 Status\n\nThis package is still in its infancy and the API is subject to change. Currently, the package can be used to train JEMs for classification. It is also possible to train pure Energy-Based Models (EBMs) for the generative task only. The package is compatible with `Flux.jl`. Work on compatibility with `MLJ.jl` (through `MLJFlux.jl`) is currently under way.\n\nWe welcome contributions and feedback at this early stage. To install the development version of the package you can run the following command:\n\n```{.julia}\nusing Pkg\nPkg.add(url=\"https://github.com/juliatrustworthyai/JointEnergyModels.jl\")\n```\n\n## 🔍 Usage Example\n\nBelow we first generate some synthetic data:\n\n::: {.cell execution_count=2}\n``` {.julia .cell-code}\nnobs=2000\nX, y = make_circles(nobs, noise=0.1, factor=0.5)\nXplot = Float32.(permutedims(matrix(X)))\nX = table(permutedims(Xplot))\nplt = scatter(Xplot[1,:], Xplot[2,:], group=y, label=\"\")\nbatch_size = Int(round(nobs/10))\ndisplay(plt)\n```\n\n::: {.cell-output .cell-output-display}\n![](index_files/figure-commonmark/cell-3-output-1.svg){}\n:::\n:::\n\n\nThe `MLJ` compatible classifier can be instantiated as follows:\n\n::: {.cell execution_count=3}\n``` {.julia .cell-code}\n𝒟x = Normal()\n𝒟y = Categorical(ones(2) ./ 2)\nsampler = ConditionalSampler(𝒟x, 𝒟y, input_size=size(Xplot)[1:end-1], batch_size=batch_size)\nclf = JointEnergyClassifier(\n    sampler;\n    builder=MLJFlux.MLP(hidden=(32, 32, 32,), σ=Flux.relu),\n    batch_size=batch_size,\n    finaliser=x -> x,\n    loss=Flux.Losses.logitcrossentropy,\n)\n```\n:::\n\n\nIt uses the `MLJFlux` package to build the model:\n\n::: {.cell execution_count=4}\n``` {.julia .cell-code}\nprintln(typeof(clf) <: MLJFlux.MLJFluxModel)\n```\n\n::: {.cell-output .cell-output-stdout}\n```\ntrue\n```\n:::\n:::\n\n\nThe model can be wrapped in data and trained using the `fit!` function:\n\n::: {.cell execution_count=5}\n``` {.julia .cell-code}\nmach = machine(clf, X, y)\nfit!(mach)\n```\n:::\n\n\nThe results are visualised below. The model has learned to discriminate between the two classes (as indicated by the contours) and to generate samples from each class (as indicated by the stars).\n\n\n\n::: {.cell execution_count=7}\n\n::: {.cell-output .cell-output-display}\n![](index_files/figure-commonmark/cell-8-output-1.svg){}\n:::\n:::\n\n\n## 🎓 References\n\n",
     "supporting": [
       "index_files"
     ],