feat: add Ops.batch

[avik-pal]

upstream issues: EnzymeAD/Enzyme-JAX#239

Example

using Reactant
using Reactant: Ops

x_ra = Reactant.to_rarray(rand(Float32, 4, 16))
y_ra = Reactant.to_rarray(rand(Float32, 16, 4))

begin
    dot6(x, y) = (Ops.multiply(x.x, y.x[1]),)
    bfn1(x, y) = Ops.batch(dot6, (; x=x), (; x=(y,)); batch_dims=((; x=2), (; x=(1,))))
    @code_hlo optimize = false bfn1(x_ra, y_ra)
end

module {
  func.func private @dot6_batch(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>, tensor<4xf32>) {
    %0 = stablehlo.multiply %arg0, %arg1 : tensor<4xf32>
    return %0, %arg0, %arg1 : tensor<4xf32>, tensor<4xf32>, tensor<4xf32>
  }
  func.func @main(%arg0: tensor<16x4xf32>, %arg1: tensor<4x16xf32>) -> (tensor<16x4xf32>, tensor<16x4xf32>, tensor<4x16xf32>) {
    %0 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<16x4xf32>) -> tensor<4x16xf32>
    %1 = stablehlo.transpose %arg1, dims = [1, 0] : (tensor<4x16xf32>) -> tensor<16x4xf32>
    %2 = stablehlo.transpose %0, dims = [1, 0] : (tensor<4x16xf32>) -> tensor<16x4xf32>
    %3 = stablehlo.transpose %1, dims = [0, 1] : (tensor<16x4xf32>) -> tensor<16x4xf32>
    %4:3 = enzyme.batch @dot6_batch(%2, %3) {batch_shape = array<i64: 16>} : (tensor<16x4xf32>, tensor<16x4xf32>) -> (tensor<16x4xf32>, tensor<16x4xf32>, tensor<16x4xf32>)
    %5 = stablehlo.transpose %4#0, dims = [1, 0] : (tensor<16x4xf32>) -> tensor<4x16xf32>
    %6 = stablehlo.transpose %5, dims = [1, 0] : (tensor<4x16xf32>) -> tensor<16x4xf32>
    %7 = stablehlo.transpose %0, dims = [1, 0] : (tensor<4x16xf32>) -> tensor<16x4xf32>
    %8 = stablehlo.transpose %1, dims = [1, 0] : (tensor<16x4xf32>) -> tensor<4x16xf32>
    return %6, %7, %8 : tensor<16x4xf32>, tensor<16x4xf32>, tensor<4x16xf32>
  }
}

TODOs:

• tests
• docs
• support linearization and delinearization

[avik-pal]

I tried not using functors 9ef1fff (#535) but it gets quite messy and I couldn't figure out how to reconstruct arguments with corrected batch dims

[avik-pal]

This is going to be a fun rebase 😓

[wsmoses]

oof yeah, the "lets speed up to_rarray" just got in (and cut down time from minutes to seconds for big cases I tested on)

[github-actions]

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

[github-actions]

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

	batchmode
	batchmode,

[github-actions]

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

	TracedRArray{unwrapped_eltype(dest),ndims(dest)},
	TracedUtils.elem_apply(bc.f, args...),
	Ops.elem_apply(bc.f, args...),
	TracedRArray{unwrapped_eltype(dest),ndims(dest)}, Ops.elem_apply(bc.f, args...)