docs: setup batching tutorial

Author: avik-pal

docs/make.jl

@@ -24,8 +24,11 @@ examples = [
 pages = [
     "Reactant.jl" => "index.md",
     "Introduction" => ["Getting Started" => "introduction/index.md"],
-    "Tutorials" =>
-        ["Overview" => "tutorials/index.md", "Profiling" => "tutorials/profiling.md"],
+    "Tutorials" => [
+        "Overview" => "tutorials/index.md",
+        "Profiling" => "tutorials/profiling.md",
+        "Batching Functions with `Reactant.Ops.batch`" => "tutorials/batching.md",
+    ],
     "API Reference" => [
         "Reactant API" => "api/api.md",
         "Ops" => "api/ops.md",
@@ -38,6 +41,11 @@ pages = [
             "Func" => "api/func.md",
             "StableHLO" => "api/stablehlo.md",
             "VHLO" => "api/vhlo.md",
+            "GPU" => "api/gpu.md",
+            "LLVM" => "api/llvm.md",
+            "NVVM" => "api/nvvm.md",
+            "TPU" => "api/tpu.md",
+            "Triton" => "api/triton.md",
         ],
         "MLIR API" => "api/mlirc.md",
         "XLA" => "api/xla.md",

docs/src/.vitepress/config.mts

@@ -56,8 +56,12 @@ export default defineConfig({
       {
         text: "Tutorials",
         items: [
-          {text: "Overview", link: "/tutorials/"},
+          { text: "Overview", link: "/tutorials/" },
           {text: "Profiling", link: "/tutorials/profiling"},
+          {
+            text: "Batching Functions with `Reactant.Ops.batch`",
+            link: "/tutorials/batching"
+          },
         ],
       },
       {
@@ -112,6 +116,10 @@ export default defineConfig({
         items: [
           { text: "Overview", link: "/tutorials/" },
           { text: "Profiling", link: "/tutorials/profiling" },
+          {
+            text: "Batching Functions with `Reactant.Ops.batch`",
+            link: "/tutorials/batching",
+          },
         ],
       },
       "/api/": {

docs/src/tutorials/batching.md

@@ -0,0 +1,3 @@
+# [Batching Functions with [`Reactant.Ops.batch`](@ref)](@id batching-tutorial)
+
+

docs/src/tutorials/index.md

@@ -1,5 +1,6 @@
 # Tutorials
 
  - [Profiling](@ref profiling).
+ - [Batching Functions with `Reactant.Ops.batch`](@ref batching-tutorial)
 
 We are currently working on adding more tutorials to Reactant!! Please check back soon!

src/Compiler.jl

@@ -776,7 +776,7 @@ function codegen_unflatten!(
         paths = (
             (
                 p for p in Reactant.TracedUtils.get_paths(result) if
-                length(p) ≥ 1 && (p[1] == :result || p[1] == :resargs)
+                length(p) > 0 && (p[1] == :result || p[1] == :resargs)
             )...,
         )
         for path in paths
@@ -846,7 +846,7 @@ function codegen_unflatten!(
         paths = (
             (
                 p for p in Reactant.TracedUtils.get_paths(result) if
-                length(p) ≥ 1 && (p[1] == :result || p[1] == :resargs || p[1] == :args)
+                length(p) > 0 && (p[1] == :result || p[1] == :resargs || p[1] == :args)
             )...,
         )
 

src/Ops.jl

@@ -2013,8 +2013,24 @@ end
 # This function assumes that the last dimension of each element is the batch dimension by
 # default. This is the standard Julia ordering for batching. We permutedims the ordering to
 # make sure the first dimension is the batch dimension when calling `batch_internal` below.
-# XXX: Mutation inside a batched function is not supported yet (need to set the results
-#      correctly)
+"""
+    batch(f, args...; batch_dims=nothing, result_dims=nothing)
+
+Map `f` over the arguments `args` along the batch dimensions `batch_dims` and return the results with the corresponding batch dimensions specified by `result_dims`. (For users
+familiar with `jax`, this operation corresponds to `jax.vmap`.)
+
+If `batch_dims` is `nothing`, we assume that the last dimension of each leaf of `args` is the batch dimension. If `result_dims` is `nothing`, we assume that the last dimension of each leaf of the returned values is the batch dimension.
+
+To avoid batching a specific leaf, pass `nothing` for the corresponding `batch_dims`.
+
+## Examples
+
+For usage examples, see the [Batching Functions with `Reactant.Ops.batch`](@ref batching-tutorial) tutorial.
+
+!!! danger
+
+    Mutation inside a batched function is not supported yet and will lead to unexpected results.
+"""
 @noinline function batch(f, args...; batch_dims=nothing, result_dims=nothing)
     batch_sizes = Int64[]
     batching_dims = if batch_dims === nothing
@@ -2060,6 +2076,8 @@ end
     end
 
     return fmap(results, result_dims) do result, dim
+        @assert dim !== nothing "Result batch dimension cannot be `nothing`"
+
         order = collect(Int64, 1:ndims(result))
         order[dim] = 1
         order[1] = dim

src/Reactant.jl

@@ -9,8 +9,6 @@ using Functors: @leaf
 using Adapt: Adapt, WrappedArray
 using GPUArraysCore: GPUArraysCore, @allowscalar, allowscalar # keep this import to allow users to do `Reactant.allowscalar(false)`
 
-using Functors: @leaf
-
 export @allowscalar # re-exported from GPUArraysCore
 
 # auxiliary types and functions

test/batching.jl

@@ -0,0 +1,2 @@
+using Reactant, Test
+

test/runtests.jl

@@ -57,6 +57,7 @@ const REACTANT_TEST_GROUP = lowercase(get(ENV, "REACTANT_TEST_GROUP", "all"))
         @safetestset "Wrapped Arrays" include("wrapped_arrays.jl")
         @safetestset "Control Flow" include("control_flow.jl")
         @safetestset "Sorting" include("sorting.jl")
+        @safetestset "Batching" include("batching.jl")
     end
 
     if REACTANT_TEST_GROUP == "all" || REACTANT_TEST_GROUP == "integration"