Update documentation

README.md

@@ -1,6 +1,9 @@
 # ParametricOperators.jl
 
-[![][license-img]][license-status] 
+[![][license-img]][license-status]
+[![Documenter](https://github.com/slimgroup/ParametricOperators.jl/actions/workflows/Documenter.yml/badge.svg)](https://github.com/slimgroup/ParametricOperators.jl/actions/workflows/Documenter.yml)
+[![TagBot](https://github.com/slimgroup/ParametricOperators.jl/actions/workflows/TagBot.yml/badge.svg)](https://github.com/slimgroup/ParametricOperators.jl/actions/workflows/TagBot.yml)
+
 <!-- [![][zenodo-img]][zenodo-status] -->
 
 ParametricOperators.jl is a Julia Language-based scientific library designed to facilitate the creation and manipulation of tensor operations involving large-scale data using Kronecker products. It provides an efficient and mathematically consistent way to express tensor programs and distribution in the context of machine learning.
@@ -16,13 +19,17 @@ ParametricOperators.jl is a Julia Language-based scientific library designed to
 
    ```julia
    julia> using Pkg
-   julia> Pkg.activate("path/to/your/project")
+   julia> Pkg.activate("path/to/your/environment")
    julia> Pkg.add("ParametricOperators")
    ```
 
 This will add `ParametricOperators.jl` as dependency to your project
 
-## Examples
+## Documentation
+
+[Detailed Documentation hosted here](https://slimgroup.github.io/ParametricOperators.jl)
+
+## Some Examples
 
 ### 1. FFT of 3D Tensor
 
@@ -55,7 +62,7 @@ Make sure to add necessary dependencies. You might also need to load a proper MP
 
 ```julia
 julia> using Pkg
-julia> Pkg.activate("path/to/your/project")
+julia> Pkg.activate("path/to/your/environment")
 julia> Pkg.add("MPI")
 julia> Pkg.add("CUDA")
 ```
@@ -111,7 +118,7 @@ Make sure to add necessary dependencies to compute the gradient
 
 ```julia
 julia> using Pkg
-julia> Pkg.activate("path/to/your/project")
+julia> Pkg.activate("path/to/your/environment")
 julia> Pkg.add("Zygote")
 ```
 
@@ -151,7 +158,7 @@ Make sure to add necessary dependencies. You might also need to load a proper MP
 
 ```julia
 julia> using Pkg
-julia> Pkg.activate("path/to/your/project")
+julia> Pkg.activate("path/to/your/environment")
 julia> Pkg.add("MPI")
 julia> Pkg.add("CUDA")
 julia> Pkg.add("Zygote")

docs/make.jl

@@ -5,6 +5,17 @@ makedocs(
     sitename = "ParametricOperators",
     format = Documenter.HTML(),
     # modules = [ParametricOperators]
+    pages=[
+        "Introduction" => "index.md",
+        "Quick Start" => "quickstart.md",
+        "Examples" => [
+            "3D FFT" => "examples/3D_FFT.md",
+            "Distributed 3D FFT" => "examples/3D_DFFT.md",
+            "3D Conv" => "examples/3D_Conv.md",
+            "Distributed 3D Conv" => "examples/3D_DConv.md",
+        ],
+        "API" => "api.md"
+    ]
 )
 
 # Automatically deploy documentation to gh-pages.

docs/src/api.md

@@ -0,0 +1,3 @@
+# ParametricOperators.jl
+
+Documentation for ParametricOperators.jl

docs/src/examples.md

@@ -0,0 +1,3 @@
+# ParametricOperators.jl
+
+Documentation for ParametricOperators.jl

docs/src/examples/3D_Conv.md

@@ -0,0 +1,39 @@
+### Parametrized Convolution on 3D Tensor
+
+Make sure to add necessary dependencies to compute the gradient
+
+```julia
+julia> using Pkg
+julia> Pkg.activate("path/to/your/environment")
+julia> Pkg.add("Zygote")
+```
+
+```julia
+using Pkg
+Pkg.activate("./path/to/your/environment")
+
+using ParametricOperators
+using Zygote
+
+T = Float32
+
+gt, gx, gy = 100, 100, 100
+
+# Define a transform along each dimension
+St = ParMatrix(T, gt, gt)
+Sx = ParMatrix(T, gx, gx)
+Sy = ParMatrix(T, gy, gy)
+
+# Create a Kronecker operator than chains together the transforms
+S = Sy ⊗ Sx ⊗ St
+
+# Parametrize our transform
+θ = init(S) |> gpu
+
+# Apply the transform on a random input
+x = rand(T, gt, gx, gy) |> gpu
+y = S(θ) * vec(x)
+
+# Compute the gradient wrt some objective of our parameters
+θ′ = gradient(θ -> sum(S(θ) * vec(x)), θ)
+```

docs/src/examples/3D_DConv.md

@@ -0,0 +1,62 @@
+### Distributed Parametrized Convolution of a 3D Tensor:
+
+Make sure to add necessary dependencies. You might also need to load a proper MPI implementation based on your hardware.
+
+```julia
+julia> using Pkg
+julia> Pkg.activate("path/to/your/environment")
+julia> Pkg.add("MPI")
+julia> Pkg.add("CUDA")
+julia> Pkg.add("Zygote")
+```
+
+Copy the following code into a `.jl` file
+```julia
+using Pkg
+Pkg.activate("./path/to/your/environment")
+
+using ParametricOperators
+using CUDA
+using MPI
+
+MPI.Init()
+
+comm = MPI.COMM_WORLD
+rank = MPI.Comm_rank(comm)
+size = MPI.Comm_size(comm)
+
+# Julia requires you to manually assign the gpus, modify to your case.
+CUDA.device!(rank % 4)
+partition = [1, 1, size]
+
+T = Float32
+
+# Define your Global Size and Data Partition
+gt, gx, gy = 100, 100, 100
+nt, nx, ny = [gt, gx, gy] .÷ partition
+
+# Define a transform along each dimension
+St = ParMatrix(T, gt, gt)
+Sx = ParMatrix(T, gx, gx)
+Sy = ParMatrix(T, gy, gy)
+
+# Create and distribute the Kronecker operator than chains together the transforms
+S = Sy ⊗ Sx ⊗ St
+S = distribute(S, partition)
+
+# Parametrize our transform
+θ = init(S) |> gpu
+
+# Apply the transform on a random input
+x = rand(T, nt, nx, ny) |> gpu
+y = S(θ) * vec(x)
+
+# Compute the gradient wrt some objective of our parameters
+θ′ = gradient(θ -> sum(S(θ) * vec(x)), θ)
+
+MPI.Finalize()
+```
+
+You can run the above by doing:
+
+`srun -n N_TASKS julia code_above.jl`

docs/src/examples/3D_DFFT.md

@@ -0,0 +1,55 @@
+### Distributed FFT of a 3D Tensor:
+
+Make sure to add necessary dependencies. You might also need to load a proper MPI implementation based on your hardware.
+
+```julia
+julia> using Pkg
+julia> Pkg.activate("path/to/your/environment")
+julia> Pkg.add("MPI")
+julia> Pkg.add("CUDA")
+```
+
+Copy the following code into a `.jl` file
+```julia
+using Pkg
+Pkg.activate("./path/to/your/environment")
+
+using ParametricOperators
+using CUDA
+using MPI
+
+MPI.Init()
+
+comm = MPI.COMM_WORLD
+rank = MPI.Comm_rank(comm)
+size = MPI.Comm_size(comm)
+
+# Julia requires you to manually assign the gpus, modify to your case.
+CUDA.device!(rank % 4)
+partition = [1, 1, size]
+
+T = Float32
+
+# Define your Global Size and Data Partition
+gt, gx, gy = 100, 100, 100
+nt, nx, ny = [gt, gx, gy] .÷ partition
+
+# Define a transform along each dimension
+Ft = ParDFT(T, gt)
+Fx = ParDFT(Complex{T}, gx)
+Fy = ParDFT(Complex{T}, gy)
+
+# Create and distribute the Kronecker operator than chains together the transforms
+F = Fy ⊗ Fx ⊗ Ft
+F = distribute(F, partition)
+
+# Apply the transform on a random input
+x = rand(T, nt, nx, ny) |> gpu
+y = F * vec(x)
+
+MPI.Finalize()
+```
+
+You can run the above by doing:
+
+`srun -n N_TASKS julia code_above.jl`

docs/src/examples/3D_FFT.md

@@ -0,0 +1,24 @@
+### FFT of 3D Tensor
+
+```julia
+using Pkg
+Pkg.activate("./path/to/your/environment")
+
+using ParametricOperators
+
+T = Float32
+
+gt, gx, gy = 100, 100, 100
+
+# Define a transform along each dimension
+Ft = ParDFT(T, gt)
+Fx = ParDFT(Complex{T}, gx)
+Fy = ParDFT(Complex{T}, gy)
+
+# Create a Kronecker operator than chains together the transforms
+F = Fy ⊗ Fx ⊗ Ft
+
+# Apply the transform on a random input
+x = rand(T, gt, gx, gy) |> gpu
+y = F * vec(x)
+```

docs/src/index.md

@@ -1,3 +1,3 @@
-# ParametricOperators.jl
+### ParametricOperators.jl
 
 Documentation for ParametricOperators.jl

docs/src/quickstart.md

@@ -0,0 +1,72 @@
+## Installation
+
+Add `ParametricOperators.jl` as a dependency to your environment.
+
+To add, either do:
+
+```julia
+julia> ]
+(v1.9) add ParametricOperators
+```
+
+OR
+
+```julia
+julia> using Pkg
+julia> Pkg.activate("path/to/your/environment")
+julia> Pkg.add("ParametricOperators")
+```
+
+## Simple Operator
+
+Make sure to include the package in your environment
+
+```julia
+using ParametricOperators
+```
+
+Lets start by defining a Matrix Operator of size `10x10`:
+
+```julia
+A = ParMatrix(10, 10)
+```
+
+Now, we parametrize our operator with some weights `θ`:
+
+```julia
+θ = init(A)
+```
+
+We can now apply our operator on some random input:
+
+```julia
+x = rand(10)
+A(θ) * x
+```
+
+## Gradient Computation
+
+!!! note "Limited AD support"
+    Current support only provided for Zygote.jl
+
+Make sure to include an AD package in your environment
+
+```julia
+using Zygote
+```
+
+Using the example above, one can find the gradient of the weights or your input w.r.t to some objective using a standard AD package:
+
+```julia
+# Gradient w.r.t weights
+θ′ = gradient(θ -> sum(A(θ) * x), θ)
+
+# Gradient w.r.t input
+x′ = gradient(x -> sum(A(θ) * x), x)
+```
+
+## Chaining Operators
+
+## Kronecker Operator
+
+## Distributing Kronecker Operator