feat: new batching opt passes (#1582)

* feat: new opt passes

* feat: new batching passes [skip ci]

* chore: bump jll

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Reactant"
 uuid = "3c362404-f566-11ee-1572-e11a4b42c853"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>", "Sergio Sánchez Ramírez <[email protected]>", "Paul Berg <[email protected]>", "Avik Pal <[email protected]>", "Mosè Giordano <[email protected]>"]
-version = "0.2.157"
+version = "0.2.158"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -100,7 +100,7 @@ PythonCall = "0.9.25"
 Random = "1.10"
 Random123 = "1.7"
 ReactantCore = "0.1.15"
-Reactant_jll = "0.0.235"
+Reactant_jll = "0.0.236"
 ScopedValues = "1.3.0"
 Scratch = "1.2"
 Sockets = "1.10"

src/Compiler.jl

@@ -907,6 +907,12 @@ function optimization_passes(
         "self_add_to_convolution_like($(Int(backend == "tpu")))",
         "self_mul_to_convolution_like($(Int(backend == "tpu")))",
         "subtract_multiply_const_to_add_mul_const",
+        "concat_insert_dim_dot_general",
+        "concat_insert_dim_gather",
+        "concat_insert_dim_iota",
+        "concat_insert_dim_reduce",
+        "concat_insert_dim_sort",
+        "concat_insert_dim_reduce_window",
     ]
 
     if !compile_options.disable_scatter_gather_optimization_passes