Format code

Author: Pangoraw

ext/ReactantNNlibExt.jl

@@ -60,15 +60,16 @@ function NNlib.conv(
 
     in1, in2 = in1.mlir_data, in2.mlir_data
     if !NNlib.flipkernel(cdims)
-        rev = Reactant.MLIR.Dialects.stablehlo.reverse(in2; dimensions=[1,0])
+        rev = Reactant.MLIR.Dialects.stablehlo.reverse(in2; dimensions=[1, 0])
         in2 = Reactant.MLIR.IR.result(rev, 1)
     end
 
     return Reactant.TracedRArray{T,output_size,4}(
         (),
         Reactant.MLIR.IR.result(
             Reactant.MLIR.Dialects.stablehlo.convolution(
-                in1, in2;
+                in1,
+                in2;
                 result_0=output_type,
                 dimension_numbers=parse(
                     Reactant.MLIR.IR.Attribute,
@@ -88,7 +89,8 @@ function NNlib.conv(
                 ),
                 rhs_dilation,
                 lhs_dilation,
-                padding, window_strides,
+                padding,
+                window_strides,
                 feature_group_count=1,
                 batch_group_count=1,
             ),

src/Reactant.jl

@@ -1174,7 +1174,7 @@ function run_pass_pipeline!(mod, pass_pipeline)
     pm = MLIR.IR.PassManager()
     opm = MLIR.IR.OpPassManager(pm)
     MLIR.IR.add_pipeline!(opm, pass_pipeline)
-    MLIR.IR.run!(pm, mod)
+    return MLIR.IR.run!(pm, mod)
 end
 
 function compile_to_module(mod, f, args; optimize=true)

test/nn.jl

@@ -75,7 +75,7 @@ mean((out2[1, :] .> 0.5) .== truth)  # accuracy 94% so far!
     W = randn(Float32, 10, 10, 3, 1)
     x = randn(Float32, 64, 64, 3, 2)
 
-    cW, cx = (W, x) .|> Reactant.ConcreteRArray
+    cW, cx = Reactant.ConcreteRArray.((W, x))
     cconv = Reactant.compile(NNlib.conv, (cx, cW))
 
     out = NNlib.conv(x, W)