doc fixes again

Author: ayushinav

src/deeponet.jl

@@ -1,6 +1,6 @@
 """
-	DeepONet(; branch = (64, 32, 32, 16), trunk = (1, 8, 8, 16),
-		branch_activation = identity, trunk_activation = identity)
+    DeepONet(; branch = (64, 32, 32, 16), trunk = (1, 8, 8, 16),
+    	branch_activation = identity, trunk_activation = identity)
 
 Constructs a DeepONet composed of Dense layers. Make sure the last node of `branch` and
 `trunk` are same.
@@ -44,25 +44,25 @@ Trunk net :
 )
 ```
 """
-function DeepONet(; branch = (64, 32, 32, 16), trunk = (1, 8, 8, 16),
-	branch_activation = identity, trunk_activation = identity)
+function DeepONet(; branch=(64, 32, 32, 16), trunk=(1, 8, 8, 16),
+        branch_activation=identity, trunk_activation=identity)
 
-	# checks for last dimension size
-	@argcheck branch[end] == trunk[end] "Branch and Trunk net must share the same amount of \
-										 nodes in the last layer. Otherwise Σᵢ bᵢⱼ tᵢₖ won't \
-										 work."
+    # checks for last dimension size
+    @argcheck branch[end]==trunk[end] "Branch and Trunk net must share the same amount of \
+            nodes in the last layer. Otherwise Σᵢ bᵢⱼ tᵢₖ won't \
+            work."
 
-	branch_net = Chain([Dense(branch[i] => branch[i+1], branch_activation)
-						for i in 1:(length(branch)-1)]...)
+    branch_net = Chain([Dense(branch[i] => branch[i + 1], branch_activation)
+                        for i in 1:(length(branch) - 1)]...)
 
-	trunk_net = Chain([Dense(trunk[i] => trunk[i+1], trunk_activation)
-					   for i in 1:(length(trunk)-1)]...)
+    trunk_net = Chain([Dense(trunk[i] => trunk[i + 1], trunk_activation)
+                       for i in 1:(length(trunk) - 1)]...)
 
-	return DeepONet(branch_net, trunk_net)
+    return DeepONet(branch_net, trunk_net)
 end
 
 """
-	DeepONet(branch, trunk)
+    DeepONet(branch, trunk)
 
 Constructs a DeepONet from a `branch` and `trunk` architectures. Make sure that both the
 nets output should have the same first dimension.
@@ -107,28 +107,28 @@ Trunk net :
 ```
 """
 function DeepONet(branch::L1, trunk::L2) where {L1, L2}
-	return @compact(; branch, trunk, dispatch = :DeepONet) do (u, y) # ::AbstractArray{<:Real, M} where {M}
-		t = trunk(y) # p x N x nb
-		b = branch(u) # p x nb
-
-		# checks for last dimension size
-		@argcheck size(t, 1) == size(b, 1) "Branch and Trunk net must share the same amount \
-											of nodes in the last layer. Otherwise Σᵢ bᵢⱼ tᵢₖ \
-											won't work."
-
-		tᵀ = permutedims(t, (2, 1, 3)) # N x p x nb
-		b_ = permutedims(reshape(b, size(b)..., 1), (1, 3, 2)) # p x 1 x nb
-		G = batched_mul(tᵀ, b_) # N x 1 X nb
-		@return dropdims(G; dims = 2)
-	end
+    return @compact(; branch, trunk, dispatch=:DeepONet) do (u, y) # ::AbstractArray{<:Real, M} where {M}
+        t = trunk(y) # p x N x nb
+        b = branch(u) # p x nb
+
+        # checks for last dimension size
+        @argcheck size(t, 1)==size(b, 1) "Branch and Trunk net must share the same amount \
+               of nodes in the last layer. Otherwise Σᵢ bᵢⱼ tᵢₖ \
+               won't work."
+
+        tᵀ = permutedims(t, (2, 1, 3)) # N x p x nb
+        b_ = permutedims(reshape(b, size(b)..., 1), (1, 3, 2)) # p x 1 x nb
+        G = batched_mul(tᵀ, b_) # N x 1 X nb
+        @return dropdims(G; dims=2)
+    end
 end
 
 function Base.show(io::IO, model::Lux.CompactLuxLayer{:DeepONet})
-	Lux._print_wrapper_model(io, "Branch net :\n", model.layers.branch)
-	print(io, "\n \n")
-	Lux._print_wrapper_model(io, "Trunk net :\n", model.layers.trunk)
+    Lux._print_wrapper_model(io, "Branch net :\n", model.layers.branch)
+    print(io, "\n \n")
+    Lux._print_wrapper_model(io, "Trunk net :\n", model.layers.trunk)
 end
 
 function Base.show(io::IO, ::MIME"text/plain", x::CompactLuxLayer{:DeepONet})
-	show(io, x)
+    show(io, x)
 end

test/deeponet_tests.jl

@@ -1,35 +1,35 @@
-@testitem "DeepONet" setup = [SharedTestSetup] begin
-	@testset "BACKEND: $(mode)" for (mode, aType, dev, ongpu) in MODES
-		rng_ = get_stable_rng()
+@testitem "DeepONet" setup=[SharedTestSetup] begin
+    @testset "BACKEND: $(mode)" for (mode, aType, dev, ongpu) in MODES
+        rng_ = get_stable_rng()
 
-		u = rand(64, 5) |> aType # sensor_points x nb
-		y = rand(1, 10, 5) |> aType # ndims x N x nb
-		out_size = (10, 5)
+        u = rand(64, 5) |> aType # sensor_points x nb
+        y = rand(1, 10, 5) |> aType # ndims x N x nb
+        out_size = (10, 5)
 
-		don_ = DeepONet(; branch = (64, 32, 32, 16), trunk = (1, 8, 8, 16))
+        don_ = DeepONet(; branch=(64, 32, 32, 16), trunk=(1, 8, 8, 16))
 
-		ps, st = Lux.setup(rng_, don_) |> dev
+        ps, st = Lux.setup(rng_, don_) |> dev
 
-		@inferred don_((u, y), ps, st)
-		@jet don_((u, y), ps, st)
+        @inferred don_((u, y), ps, st)
+        @jet don_((u, y), ps, st)
 
-		pred = first(don_((u, y), ps, st))
-		@test size(pred) == out_size
+        pred = first(don_((u, y), ps, st))
+        @test size(pred) == out_size
 
-		don_ = DeepONet(Chain(Dense(64 => 32), Dense(32 => 32), Dense(32 => 16)),
-			Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16)))
+        don_ = DeepONet(Chain(Dense(64 => 32), Dense(32 => 32), Dense(32 => 16)),
+            Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16)))
 
-		ps, st = Lux.setup(rng_, don_) |> dev
+        ps, st = Lux.setup(rng_, don_) |> dev
 
-		@inferred don_((u, y), ps, st)
-		@jet don_((u, y), ps, st)
+        @inferred don_((u, y), ps, st)
+        @jet don_((u, y), ps, st)
 
-		pred = first(don_((u, y), ps, st))
-		@test size(pred) == out_size
+        pred = first(don_((u, y), ps, st))
+        @test size(pred) == out_size
 
-		don_ = DeepONet(Chain(Dense(64 => 32), Dense(32 => 32), Dense(32 => 20)),
-			Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16)))
-		ps, st = Lux.setup(rng_, don_) |> dev
-		@test_throws ArgumentError don_((u, y), ps, st)
-	end
+        don_ = DeepONet(Chain(Dense(64 => 32), Dense(32 => 32), Dense(32 => 20)),
+            Chain(Dense(1 => 8), Dense(8 => 8), Dense(8 => 16)))
+        ps, st = Lux.setup(rng_, don_) |> dev
+        @test_throws ArgumentError don_((u, y), ps, st)
+    end
 end