Fix bug rewriting sum with dims kwarg

Project.toml

@@ -1,7 +1,7 @@
 name = "MutableArithmetics"
 uuid = "d8a4904e-b15c-11e9-3269-09a3773c0cb0"
 authors = ["Gilles Peiffer", "Benoît Legat", "Sascha Timme"]
-version = "1.3.2"
+version = "1.3.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/rewrite_generic.jl

@@ -51,6 +51,10 @@ function _is_parameters(expr)
  return Meta.isexpr(expr, :call, 3) && Meta.isexpr(expr.args[2], :parameters)
 end
 
+function _is_kwarg(expr, kwarg::Symbol)
+ return Meta.isexpr(expr, :kw) && expr.args[1] == kwarg
+end
+
 """
  _rewrite_generic(stack::Expr, expr::Expr)
 
@@ -78,14 +82,20 @@ function _rewrite_generic(stack::Expr, expr::Expr)
  # come in two forms: `sum(i for i=I, j=J)` or `sum(i for i=I for j=J)`.
  # The latter is a `:flatten` expression and needs additional handling,
  # but we delay this complexity for _rewrite_generic_generator.
- if Meta.isexpr(expr.args[2], :parameters, 1) &&
- Meta.isexpr(expr.args[2].args[1], :kw, 2) &&
- expr.args[2].args[1].args[1] == :init
- # sum(iter ; init) form!
- root = gensym()
- init, _ = _rewrite_generic(stack, expr.args[2].args[1].args[2])
- push!(stack.args, :($root = $init))
- return _rewrite_generic_generator(stack, :+, expr.args[3], root)
+ if Meta.isexpr(expr.args[2], :parameters)
+ # The summation has keyword arguments. We can deal with `init`, but
+ # not any of the others.
+ p = expr.args[2]
+ if length(p.args) == 1 && _is_kwarg(p.args[1], :init)
+ # sum(iter ; init) form!
+ root = gensym()
+ init, _ = _rewrite_generic(stack, p.args[1].args[2])
+ push!(stack.args, :($root = $init))
+ return _rewrite_generic_generator(stack, :+, expr.args[3], root)
+ else
+ # We don't know how to deal with this
+ return esc(expr), false
+ end
  else
  # Summations use :+ as the reduction operator.
  init_expr = expr.args[2].args[end]

test/rewrite_generic.jl

@@ -344,6 +344,59 @@ function test_rewrite_expression()
  return
 end
 
+function test_rewrite_generic_sum_dims()
+ x = [1 2; 3 4]
+ @test ==(
+ MA.@rewrite(sum(x; dims = 1), move_factors_into_sums = false),
+ [4 6],
+ )
+ @test ==(
+ MA.@rewrite(sum(x; dims = 2), move_factors_into_sums = false),
+ [3; 7;;],
+ )
+ @test ==(
+ MA.@rewrite(sum(x; dims = 1, init = 0), move_factors_into_sums = false),
+ [4 6],
+ )
+ @test ==(
+ MA.@rewrite(sum(x; dims = 2, init = 0), move_factors_into_sums = false),
+ [3; 7;;],
+ )
+ @test ==(
+ MA.@rewrite(sum(x; init = 0, dims = 1), move_factors_into_sums = false),
+ [4 6],
+ )
+ @test ==(
+ MA.@rewrite(sum(x; init = 0, dims = 2), move_factors_into_sums = false),
+ [3; 7;;],
+ )
+ @test ==(
+ MA.@rewrite(sum(x, dims = 1), move_factors_into_sums = false),
+ [4 6],
+ )
+ @test ==(
+ MA.@rewrite(sum(x, dims = 2), move_factors_into_sums = false),
+ [3; 7;;],
+ )
+ @test ==(
+ MA.@rewrite(sum(x, dims = 1, init = 0), move_factors_into_sums = false),
+ [4 6],
+ )
+ @test ==(
+ MA.@rewrite(sum(x, dims = 2, init = 0), move_factors_into_sums = false),
+ [3; 7;;],
+ )
+ @test ==(
+ MA.@rewrite(sum(x, init = 0, dims = 1), move_factors_into_sums = false),
+ [4 6],
+ )
+ @test ==(
+ MA.@rewrite(sum(x, init = 0, dims = 2), move_factors_into_sums = false),
+ [3; 7;;],
+ )
+ return
+end
+
 end # module
 
 TestRewriteGeneric.runtests()