LuxCore.jl: Extremely light dependency for Lux Compatibility (#191)

* Split out LuxCore.jl * Split out LuxCore.jl * Ordering of packages
LuxDL · Nov 18, 2022 · 7249069 · 7249069 · avik-pal · Nov 18, 2022
1 parent 1be3bf7
commit 7249069
Show file tree

Hide file tree

Showing 29 changed files with 305 additions and 256 deletions.
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -18,14 +18,14 @@ jobs:
       fail-fast: false
       matrix:
         group:
-          - Lux       # Core Framework
-          - Boltz     # Prebuilt Models using Lux
-          - LuxLib    # Backend of Lux
-          - Flux2Lux  # Flux2Lux Converter
+          - Lux # Core Framework
+          - Boltz # Prebuilt Models using Lux
+          - LuxLib # Backend of Lux
+          - Flux2Lux # Flux2Lux Converter
+          - LuxCore # Avoid the heavy Lux Dependency
         version:
-          - '1.6'      # JET tests are disabled on 1.6
-          - '1.7'
-          - '1.8'
+          - "1.6"
+          - "1.8"
     steps:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1
@@ -41,13 +41,16 @@ jobs:
             ${{ runner.os }}-test-${{ env.cache-name }}-
             ${{ runner.os }}-test-
             ${{ runner.os }}-
-      - uses: julia-actions/julia-buildpkg@v1
+      - name: Build Package
+        run: |
+          julia --project=. -e 'using Pkg; Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/LuxCore"))); Pkg.instantiate()'
       - uses: julia-actions/julia-runtest@v1
         env:
           GROUP: ${{ matrix.group }}
+          OVERRIDE_INTER_DEPENDENCIES: "false"
       - uses: julia-actions/julia-processcoverage@v1
         with:
-          directories: src,lib/Boltz/src,lib/LuxLib/src,lib/Flux2Lux/src
+          directories: src,lib/Boltz/src,lib/LuxLib/src,lib/Flux2Lux/src,lib/LuxCore/src
       - uses: codecov/codecov-action@v3
         with:
           files: lcov.info

diff --git a/.github/workflows/CINightly.yml b/.github/workflows/CINightly.yml
@@ -22,6 +22,7 @@ jobs:
           - Boltz     # Prebuilt Models using Lux
           - LuxLib    # Backend of Lux
           - Flux2Lux  # Flux2Lux Converter
+          - LuxCore   # Avoid the heavy Lux Dependency
         version:
           - 'nightly'  # merge even if tests fail
     steps:
@@ -39,13 +40,16 @@ jobs:
             ${{ runner.os }}-test-${{ env.cache-name }}-
             ${{ runner.os }}-test-
             ${{ runner.os }}-
-      - uses: julia-actions/julia-buildpkg@v1
+      - name: Build Package
+        run: |
+          julia --project=. -e 'using Pkg; Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/LuxCore"))); Pkg.instantiate()'
       - uses: julia-actions/julia-runtest@v1
         env:
           GROUP: ${{ matrix.group }}
+          OVERRIDE_INTER_DEPENDENCIES: "false"
       - uses: julia-actions/julia-processcoverage@v1
         with:
-          directories: src,lib/Boltz/src,lib/LuxLib/src,lib/Flux2Lux/src
+          directories: src,lib/Boltz/src,lib/LuxLib/src,lib/Flux2Lux/src,lib/LuxCore/src
       - uses: codecov/codecov-action@v3
         with:
           files: lcov.info

diff --git a/.github/workflows/CompatHelper.yml b/.github/workflows/CompatHelper.yml
@@ -27,7 +27,8 @@ jobs:
       - name: "Run CompatHelper"
         run: |
           import CompatHelper
-          CompatHelper.main(; subdirs=["", "examples", "lib/Boltz", "lib/LuxLib", "lib/Flux2Lux"])
+          CompatHelper.main(; subdirs=["", "examples", "examples/ImageNet", "lib/Boltz",
+                                       "lib/LuxLib", "lib/Flux2Lux"])
         shell: julia --color=yes {0}
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}

diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml
@@ -4,7 +4,7 @@ on:
   push:
     branches:
       - main
-    tags: '*'
+    tags: "*"
   pull_request:
 concurrency:
   # Skip intermediate builds: always.
@@ -18,7 +18,7 @@ jobs:
       - uses: actions/checkout@v2
       - uses: julia-actions/setup-julia@v1
         with:
-          version: '1.8'
+          version: "1.8"
       - uses: actions/cache@v1
         env:
           cache-name: cache-artifacts
@@ -30,9 +30,9 @@ jobs:
             ${{ runner.os }}-test-
             ${{ runner.os }}-
       - name: Install documentation dependencies
-        run: julia --project=docs -e 'using Pkg; Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/Flux2Lux"))); Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/LuxLib"))); Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
+        run: julia --project=docs -e 'using Pkg; Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/LuxCore"))); Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/LuxLib"))); Pkg.develop(PackageSpec(path=pwd())); Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/Flux2Lux"))); Pkg.instantiate()'
       - name: Install examples dependencies
-        run: julia --project=examples -e 'using Pkg; Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
+        run: julia --project=examples -e 'using Pkg; Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/LuxCore"))); Pkg.develop(PackageSpec(path=pwd())); Pkg.instantiate()'
       - name: Build and deploy
         env:
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} # For authentication with GitHub Actions token

diff --git a/.github/workflows/Downstream.yml b/.github/workflows/Downstream.yml
@@ -20,7 +20,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        julia-version: [1.7]
+        julia-version: [1.8]
         os: [ubuntu-latest]
         package:
           - { user: SciML, repo: DiffEqFlux.jl, group: BasicNeuralDE }
@@ -33,7 +33,9 @@ jobs:
         with:
           version: ${{ matrix.julia-version }}
           arch: x64
-      - uses: julia-actions/julia-buildpkg@latest
+      - name: Build Package
+        run: |
+          julia --project=. -e 'using Pkg; Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/LuxCore"))); Pkg.instantiate()'
       - name: Clone Downstream
         uses: actions/checkout@v2
         with:
@@ -57,6 +59,6 @@ jobs:
             exit(0)  # Exit immediately, as a success
           end
       - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v2
+      - uses: codecov/codecov-action@v3
         with:
           files: lcov.info
diff --git a/.github/workflows/FormatCheck.yml b/.github/workflows/FormatCheck.yml
@@ -13,7 +13,7 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        julia-version: [1.7]
+        julia-version: [1.8]
         julia-arch: [x86]
         os: [ubuntu-latest]
     steps:

diff --git a/.github/workflows/Invalidations.yml b/.github/workflows/Invalidations.yml
@@ -20,14 +20,20 @@ jobs:
       with:
         version: '1'
     - uses: actions/checkout@v3
-    - uses: julia-actions/julia-buildpkg@v1
+    # - uses: julia-actions/julia-buildpkg@v1
+    - name: Build Package
+      run: |
+        julia --project=. -e 'using Pkg; Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/LuxCore"))); Pkg.instantiate()'
     - uses: julia-actions/julia-invalidations@v1
       id: invs_pr
 
     - uses: actions/checkout@v3
       with:
         ref: ${{ github.event.repository.default_branch }}
-    - uses: julia-actions/julia-buildpkg@v1
+    # - uses: julia-actions/julia-buildpkg@v1
+    - name: Build Package
+      run: |
+        julia --project=. -e 'using Pkg; Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/LuxCore"))); Pkg.instantiate()'
     - uses: julia-actions/julia-invalidations@v1
       id: invs_default
 

diff --git a/Project.toml b/Project.toml
@@ -11,6 +11,7 @@ ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b"
 Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
@@ -31,6 +32,7 @@ ChainRulesCore = "1"
 ComponentArrays = "0.13"
 FillArrays = "0.13"
 Functors = "0.2, 0.3"
+LuxCore = "0.1"
 LuxLib = "0.1.7"
 NNlib = "0.8"
 Optimisers = "0.2"

diff --git a/docs/Project.toml b/docs/Project.toml
@@ -3,6 +3,7 @@ Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterMarkdown = "997ab1e6-3595-5248-9280-8efb232c3433"
 Flux2Lux = "ab51a4a6-c8c3-4b1f-af31-4b52a21037df"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
+LuxCore = "bb33d45b-7691-41d6-9220-0943567d0623"
 Lux = "b2108857-7c20-44ae-9111-449ecde12c47"
 LuxLib = "82251201-b29d-42c6-8e01-566dec8acb11"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"

diff --git a/docs/make.jl b/docs/make.jl
@@ -1,10 +1,20 @@
-using Documenter, DocumenterMarkdown, Flux2Lux, Lux, LuxLib, Pkg
+using Documenter, DocumenterMarkdown, Flux2Lux, LuxCore, Lux, LuxLib, Pkg
+
+function _setup_subdir_pkgs_index_file(subpkg)
+    src_file = joinpath(dirname(@__DIR__), "lib", subpkg, "README.md")
+    dst_file = joinpath(dirname(@__DIR__), "docs/src/lib", subpkg, "index.md")
+    rm(dst_file; force=true)
+    cp(src_file, dst_file)
+    return
+end
+
+_setup_subdir_pkgs_index_file.(["Boltz", "LuxLib", "Flux2Lux", "LuxCore"])
 
 deployconfig = Documenter.auto_detect_deploy_system()
 Documenter.post_status(deployconfig; type="pending", repo="github.com/avik-pal/Lux.jl.git")
 
 makedocs(; sitename="Lux", authors="Avik Pal et al.", clean=true, doctest=true,
-         modules=[Flux2Lux, Lux, LuxLib],
+         modules=[Flux2Lux, Lux, LuxLib, LuxCore],
          strict=[
              :doctest,
              :linkcheck,

diff --git a/docs/mkdocs.yml b/docs/mkdocs.yml
@@ -123,6 +123,10 @@ nav:
       - "LuxLib":
         - "Introduction": "lib/LuxLib/index.md"
         - "API Reference": "lib/LuxLib/api.md"
+      - "LuxCore":
+        - "Introduction": "lib/LuxCore/index.md"
+        - "API Reference": "lib/LuxCore/api.md"
   - "Development Documentation":
       - "Style Guide": "devdocs/style_guide.md"
       - "Layer Implementation": "devdocs/layer_implementation.md"
+      - "Sub Packages": "devdocs/subpackages.md"
diff --git a/docs/src/api/core.md b/docs/src/api/core.md
@@ -2,39 +2,7 @@
 CurrentModule = Lux
 ```
 
-## Abstract Types
-
-```@docs
-Lux.AbstractExplicitLayer
-Lux.AbstractExplicitContainerLayer
-```
-
-## General
-
-```@docs
-Lux.apply
-Lux.setup
-```
-
-## Parameters
-
-```@docs
-Lux.initialparameters
-Lux.parameterlength
-```
-
-## States
-
-```@docs
-Lux.initialstates
-Lux.statelength
-Lux.testmode
-Lux.trainmode
-Lux.update_state
-```
-
-## Index
-
-```@index
-Pages = ["core.md"]
-```
+The documentation for this page has been moved to [LuxCore.jl](../lib/LuxCore/api.md).
+However, this the functionality has in **no way been deprecated** and there is no plan
+to deprecated access to these functionalities in the future. The change was made merely to
+allow lighter dependencies for users extending `Lux.jl`.
diff --git a/docs/src/devdocs/layer_implementation.md b/docs/src/devdocs/layer_implementation.md
@@ -41,28 +41,3 @@ varying how calls are made at different timesteps.
 1. `reset`ing the hidden-state and memory is slightly tricky.
    1. One way would be to store a `initial_hidden_state` and `initial_memory` in the state
       alongside the `hidden_state` and `memory`.
-
-
-### RNN Blocks
-
-!!! note
-    This is currently unimplemented
-
-An example implementation would be
-
-```julia
-struct RNN{R} <: Lux.AbstractExplicitContainerLayer{(:recurrent_cell,)}
-    recurrent_cell::R
-end
-
-function (l::RNN)(x::AbstractArray{T,3}, ps::NamedTuple, st::NamedTuple) where {T}
-    x_init, x_rest = Iterators.peel(eachslice(x; dims=2))
-    (y, carry), st = l.recurrent_cell(x_init, ps, st)
-    for x in x_rest
-        (y, carry), st = l.recurrent_cell((x, carry), ps, st)
-    end
-    return y, st
-end
-```
-
-We enforce the inputs to be of the format `in_dims × sequence_length × batch_size`.
diff --git a/docs/src/devdocs/style_guide.md b/docs/src/devdocs/style_guide.md
@@ -27,12 +27,6 @@ We do have automatic formatter, which opens PR after fixing common style issues,
 * No avoiding multiply symbol -- so `2x` is invalid instead do it like other languages
   `2 * x`.
 
-## Unicode Characters
-
-* No use of unicode characters is allowed.
-* The only exception is when defining DSLs. In this particular case, how to type the unicode
-  must be properly documented.
-
 ## Testing
 
 !!! note

diff --git a/docs/src/devdocs/subpackages.md b/docs/src/devdocs/subpackages.md
@@ -0,0 +1,59 @@
+# SubPackages
+
+`Lux.jl` operates somewhat like a monorepo having a weird inter-dependency structure. So
+adding new subpackages in `lib` can be somewhat complicated. Here are the guidelines that
+need to be followed:
+
+## Package Structure
+
+  * Each subpackage should be in its own directory in `lib`.
+  * Don't have a `docs` directory (see the [Documentation Section](#documentation) for
+    details).
+  * Add a `LICENSE` file (needed to register the package independently).
+
+## Workflows
+
+  * All workflows should go in the `.github/workflows` directory (in the project root).
+  * For `CI.yml` and `CINightly.yml`
+    - add the project name to `group` matrix.
+    - Under `directories` for `julia-actions/julia-processcoverage@v1` add
+      `lib/<project name>/src`.
+  * For `CompatHelper.yml` add `lib/<project name>` to the list of `subdirs`.
+
+## Documentation
+
+  * Create a directory for the package: `docs/src/lib/<project name>`.
+  * Optionally, if you want the `index.md` page for your subpackage to be same as the
+    `README.md` file, add the package name to `_setup_subdir_pkgs_index_file.([...])` in
+    `docs/make.jl`.
+  * For generating documentation (say from docstrings) for your package, you need to update
+    the documentation pipeline:
+    - In `.github/workflows/Documentation.yml` under `Install documentation dependencies`
+      install the package using
+      `Pkg.develop(PackageSpec(path=joinpath(pwd(), "lib/<project name>")))`.
+    - Add a dependency in `docs/Project.toml` (don't add compat entries).
+    - Add `using <Project Name>` in `docs/make.jl`.
+    - Add the package name to `modules` in `docs/make.jl`.
+  * For every new page that you have added (including `index.md` if using `README.md` file)
+    update `docs/mkdocs.yml` `nav` field.
+
+## Testing
+
+  * For testing, always use `Project.toml` in the `lib/<project name>/test` directory.
+  * Write the tests as you would for a normal package.
+  * In `test/runtests.jl` add the package name to `groups` if `GROUP == "All"`.
+  * Next list any cross-dependency. When CI is run, it uses the local version of the package
+    instead of the registered version.
+    - For example, `Lux` depends on `LuxLib` so `"Lux" => [_get_lib_path("LuxLib")]`
+    - `Boltz` depends on both `Lux` and `LuxLib` (via `Lux`) so
+      `"Boltz" => [_get_lib_path("LuxLib"), dirname(@__DIR__)]`
+    - If there are no cross-dependencies, remember to add an empty vector.
+
+## Registration
+
+Registration is simply, just run `@JuliaRegistrator register subdir="lib/<project name>"`
+
+## Code Coverage
+
+If you have performed all the steps correctly the code coverage for the subpackage will
+be available under the flag `<project name>`.
diff --git a/docs/src/lib/Boltz/.gitkeep b/docs/src/lib/Boltz/.gitkeep