Despecialize Ops.dot_general (#1318)

* Despecialize `Ops.dot_general`

* Remove import

* Remove einsum tests

* Format code

Author: mofeing

src/Ops.jl

@@ -794,9 +794,9 @@ end
 #     return TracedRArray{T,N}((), res, size(lhs))
 # end
 
-@noinline function dot_general(
-    lhs::TracedRArray{T1},
-    rhs::TracedRArray{T2};
+Base.@nospecializeinfer @noinline function dot_general(
+    @nospecialize(lhs::TracedRArray{T1}),
+    @nospecialize(rhs::TracedRArray{T2});
     contracting_dimensions,
     batching_dimensions=(Int[], Int[]),
     precision_config=Reactant.DOT_GENERAL_PRECISION[],
@@ -920,63 +920,6 @@ end
     return TracedRArray{resT,length(ressize)}((), res, ressize)
 end
 
-@noinline function einsum(
-    lhs::TracedRArray{T},
-    rhs::TracedRArray{T};
-    equation::String,
-    location=mlir_stacktrace("einsum", @__FILE__, @__LINE__),
-) where {T}
-    Base.depwarn(
-        "`stablehlo.einsum` is on deprecation process; use `dot_general` instead", :einsum
-    )
-    ins, ic = split(equation, "->")
-    ia, ib = split(ins, ",")
-
-    sizea = Dict(c => d for (c, d) in zip(ia, size(lhs)))
-    sizeb = Dict(c => d for (c, d) in zip(ib, size(rhs)))
-    sizes = mergewith(sizea, sizeb) do da, db
-        da == db ? da : error("Invalid dimensions in einsum equation")
-    end
-
-    rsize = Tuple(sizes[i] for i in ic)
-    result_0 = mlir_type(TracedRArray{T,length(ic)}, rsize)
-
-    res = MLIR.IR.result(
-        stablehlo.einsum(
-            lhs.mlir_data,
-            rhs.mlir_data;
-            result_0,
-            einsum_config=MLIR.IR.Attribute(equation),
-            location,
-        ),
-    )
-    return TracedRArray{T,length(rsize)}((), res, rsize)
-end
-
-# function unary_einsum(
-#     x::TracedRArray{T};
-#     equation::String,
-#     location=mlir_stacktrace(
-#         "unary_einsum", @__FILE__, @__LINE__
-#     ),
-# ) where {T}
-#     ia, ic = split(equation, "->")
-#     sizes = Dict(c => d for (c, d) in zip(ia, size(x)))
-#     rsize = Tuple(sizes[i] for i in ic)
-#     result_0 = mlir_type(TracedRArray{T,length(ic)}, rsize)
-
-#     res = MLIR.IR.result(
-#         stablehlo.unary_einsum(
-#             x.mlir_data; result_0, einsum_config=MLIR.IR.Attribute(equation), location
-#         ),
-#     )
-#     if length(rsize) == 0
-#         return TracedRNumber{T}((), res)
-#     else
-#         return TracedRArray{T,length(rsize)}((), res, rsize)
-#     end
-# end
-
 # parallel ops
 @noinline function partition_id(;
     location=mlir_stacktrace("partition_id", @__FILE__, @__LINE__)

test/ops.jl

@@ -237,37 +237,6 @@ end
     @test Array(a)' * Array(b) == @jit f1(a, b)
 end
 
-@testset "einsum" begin
-    f1(a, b) = Ops.einsum(a, b; equation="i,i->i")
-    f2(a, b) = Ops.einsum(a, b; equation="i,j->ij")
-    f3(a, b) = Ops.einsum(a, b; equation="ij,ij->ij")
-    f4(a, b) = Ops.einsum(a, b; equation="ik,kj->ij")
-
-    for (a, b) in [
-        (Reactant.to_rarray([1, 2, 3, 4]), Reactant.to_rarray([5, 6, -7, -8])),
-        (
-            Reactant.to_rarray([1.0, 2.0, 3.0, 4.0]),
-            Reactant.to_rarray([5.0, 6.0, -7.0, -8.0]),
-        ),
-        (
-            Reactant.to_rarray([1.0 + 1im, 2.0 + 2im, 3.0 - 3im, 4.0 - 4im]),
-            Reactant.to_rarray([5.0 + 5im, 6.0 + 6im, -7.0 - 7im, -8.0 - 8im]),
-        ),
-    ]
-        @test a .* b ≈
-            @test_warn r"`stablehlo.einsum` is on deprecation process" @jit f1(a, b)
-        @test reshape(kron(Array(b), Array(a)), 4, 4) ≈
-            @test_warn r"`stablehlo.einsum` is on deprecation process" @jit f2(a, b)
-
-        x = ConcreteRArray(reshape(a, (2, 2)))
-        y = ConcreteRArray(reshape(b, (2, 2)))
-        @test x .* y ≈
-            @test_warn r"`stablehlo.einsum` is on deprecation process" @jit f3(x, y)
-        @test Array(x) * Array(y) ≈
-            @test_warn r"`stablehlo.einsum` is on deprecation process" @jit f4(x, y)
-    end
-end
-
 @testset "exponential" begin
     x = Reactant.to_rarray([1.0, 2.0, 3.0, 4.0])
     @test exp.(Array(x)) ≈ @jit Ops.exponential(x)