Use |> for moving data to devices (#1559)

* Use `|>` for moving data to devices

It's more visually suggestive of the data being moved.
Also, in an assignment expression, the main data piece/function call is
closer to the equal sign rather than the device call.

* Run JuliaFormatter over examples/

Author: abhro

examples/ConvolutionalVAE/main.jl

@@ -256,14 +256,14 @@ function main(;
     Random.seed!(rng, seed)
 
     cvae = CVAE(rng; num_latent_dims, image_shape=(image_size..., 1), max_num_filters)
-    ps, st = xdev(Lux.setup(rng, cvae))
+    ps, st = Lux.setup(rng, cvae) |> xdev
 
     z = xdev(randn(Float32, num_latent_dims, num_samples))
     decode_compiled = @compile decode(cvae, z, ps, Lux.testmode(st))
-    x = xdev(randn(Float32, image_size..., 1, batchsize))
+    x = randn(Float32, image_size..., 1, batchsize) |> xdev
     cvae_compiled = @compile cvae(x, ps, Lux.testmode(st))
 
-    train_dataloader = xdev(loadmnist(batchsize, image_size))
+    train_dataloader = loadmnist(batchsize, image_size) |> xdev
 
     opt = AdamW(; eta=learning_rate, lambda=weight_decay)
 

examples/DDIM/main.jl

@@ -501,7 +501,7 @@ function Base.getindex(ds::FlowersDataset, idxs)
     imgs = Array{Float32,4}(undef, ds.image_size..., 3, length(idxs))
     tforeach(1:length(idxs)) do i
         img = Image(load(ds.image_files[idxs[i]]))
-        copyto!(view(imgs, :, :, :, i), itemdata(apply(ds.transform, img)))
+        return copyto!(view(imgs, :, :, :, i), itemdata(apply(ds.transform, img)))
     end
     return imgs
 end

examples/GCN_Cora/main.jl

@@ -91,10 +91,10 @@ function main(;
     rng = Random.default_rng()
     Random.seed!(rng, 0)
 
-    features, targets, adj, (train_idx, val_idx, test_idx) = xdev(loadcora())
+    features, targets, adj, (train_idx, val_idx, test_idx) = loadcora() |> xdev
 
     gcn = GCN(size(features, 1), hidden_dim, size(targets, 1); nb_layers, dropout, use_bias)
-    ps, st = xdev(Lux.setup(rng, gcn))
+    ps, st = Lux.setup(rng, gcn) |> xdev
     opt = iszero(weight_decay) ? Adam(lr) : AdamW(; eta=lr, lambda=weight_decay)
 
     train_state = Training.TrainState(gcn, ps, st, opt)

examples/HyperNet/main.jl

@@ -97,8 +97,9 @@ function accuracy(model, ps, st, dataloader, data_idx)
     cdev = cpu_device()
     st = Lux.testmode(st)
     for (x, y) in dataloader
-        target_class = onecold(cdev(y))
-        predicted_class = onecold(cdev(first(model((data_idx, x), ps, st))))
+        ŷ, _ = model((data_idx, x), ps, st)
+        target_class = y |> cdev |> onecold
+        predicted_class = ŷ |> cdev |> onecold
         total_correct += sum(target_class .== predicted_class)
         total += length(target_class)
     end
@@ -111,10 +112,10 @@ function train()
     dev = reactant_device(; force=true)
 
     model = create_model()
-    dataloaders = dev(load_datasets())
+    dataloaders = load_datasets() |> dev
 
     Random.seed!(1234)
-    ps, st = dev(Lux.setup(Random.default_rng(), model))
+    ps, st = Lux.setup(Random.default_rng(), model) |> dev
 
     train_state = Training.TrainState(model, ps, st, Adam(0.0003f0))
 

examples/ImageNet/main.jl

@@ -226,7 +226,7 @@ function construct_model(;
     rng::AbstractRNG, model_name::String, model_args, pretrained::Bool=false
 )
     model = getproperty(Vision, Symbol(model_name))(model_args...; pretrained)
-    ps, st = gdev(Lux.setup(rng, model))
+    ps, st = Lux.setup(rng, model) |> gdev
 
     sensible_println("=> model `$(model_name)` created.")
     pretrained && sensible_println("==> using pre-trained model`")
@@ -549,7 +549,7 @@ Comonicon.@main function main(;
 
     ckpt = load_checkpoint(rpath)
     if !isnothing(ckpt)
-        ps, st = gdev((ckpt.ps, ckpt.st))
+        ps, st = (ckpt.ps, ckpt.st) |> gdev
         initial_step = ckpt.step
         sensible_println("=> training started from $(initial_step)")
     else

examples/NeuralODE/main.jl

@@ -128,8 +128,11 @@ function create_model(
     Random.seed!(rng, 0)
 
     ps, st = Lux.setup(rng, model)
-    ps = dev((use_named_tuple ? ps : ComponentArray(ps)))
-    st = dev(st)
+    if !use_named_tuple
+        ps = ComponentArray(ps)
+    end
+    ps = ps |> dev
+    st = st |> dev
 
     return model, ps, st
 end
@@ -239,7 +242,7 @@ model, ps, st = create_model(NeuralODE)
 
 model_stateful, ps_stateful, st_stateful = create_model(StatefulNeuralODE)
 
-x = gpu_device()(ones(Float32, 28, 28, 1, 3));
+x = ones(Float32, 28, 28, 1, 3) |> gpu_device();
 
 # NeuralODE is not type stable due to the boxing of `st`
 

examples/OptimizationIntegration/main.jl

@@ -78,7 +78,7 @@ Base.length(t::TimeWrapper) = length(t.t)
 
 Base.getindex(t::TimeWrapper, i) = TimeWrapper(t.t[i])
 
-dataloader = gdev(DataLoader((ode_data, TimeWrapper(t)); batchsize=8))
+dataloader = DataLoader((ode_data, TimeWrapper(t)); batchsize=8) |> gdev
 nothing #hide
 
 # ## Training the model
@@ -101,8 +101,8 @@ function train_model(dataloader)
     model = Chain(Dense(2, 32, tanh), Dense(32, 32, tanh), Dense(32, 2))
     ps, st = Lux.setup(Random.default_rng(), model)
 
-    ps_ca = gdev(ComponentArray(ps))
-    st = gdev(st)
+    ps_ca = ComponentArray(ps) |> gdev
+    st = st |> gdev
 
     function callback(state, l)
         if state.iter == 1 || state.iter % 25 == 0
@@ -153,7 +153,7 @@ nothing #hide
 dudt(u, p, t) = trained_model(u, p)
 prob = ODEProblem(dudt, gdev(u0), (tspan[1], tspan[2]), trained_model.ps)
 sol = solve(prob, Tsit5(); saveat=t)
-pred = cdev(convert(AbstractArray, sol))
+pred = convert(AbstractArray, sol) |> cdev
 
 begin
     fig = Figure()

examples/PolynomialFitting/main.jl

@@ -69,7 +69,7 @@ const loss_function = MSELoss()
 const cdev = cpu_device()
 const xdev = reactant_device()
 
-ps, st = xdev(Lux.setup(rng, model))
+ps, st = Lux.setup(rng, model) |> xdev
 
 # ## Training
 
@@ -104,11 +104,10 @@ forward_pass = @compile Lux.apply(
     tstate.model, xdev(x), tstate.parameters, Lux.testmode(tstate.states)
 )
 
-y_pred = cdev(
-    first(
-        forward_pass(tstate.model, xdev(x), tstate.parameters, Lux.testmode(tstate.states))
-    ),
-)
+y_pred =
+    forward_pass(tstate.model, xdev(x), tstate.parameters, Lux.testmode(tstate.states)) |>
+    first |>
+    cdev
 nothing #hide
 
 # Let's plot the results

examples/RealNVP/main.jl

@@ -223,12 +223,13 @@ function main(;
     rng = Random.default_rng()
     Random.seed!(rng, 0)
 
-    dataloader = Iterators.cycle(
-        xdev(load_moons_dataloader(rng, Float32, n_train_samples; batchsize, noise))
-    )
+    dataloader =
+        load_moons_dataloader(rng, Float32, n_train_samples; batchsize, noise) |>
+        xdev |>
+        Iterators.cycle
 
     model = RealNVP(; n_transforms, dist_dims=2, hidden_dims, n_layers)
-    ps, st = xdev(Lux.setup(rng, model))
+    ps, st = Lux.setup(rng, model) |> xdev
     opt = Adam(lr)
 
     train_state = Training.TrainState(model, ps, st, opt)

examples/SimpleChains/main.jl

@@ -74,8 +74,8 @@ end
 
 # ## Define the Training Loop
 function train(model, dev=cpu_device(); rng=Random.default_rng(), kwargs...)
-    train_dataloader, test_dataloader = dev(loadmnist(128, 0.9))
-    ps, st = dev(Lux.setup(rng, model))
+    train_dataloader, test_dataloader = loadmnist(128, 0.9) |> dev
+    ps, st = Lux.setup(rng, model) |> dev
 
     vjp = dev isa ReactantDevice ? AutoEnzyme() : AutoZygote()