Skip to content

Commit d56ae72

Browse files
SophieL1SophieL1
authored
Merge branch 'main' into bl/fix_parse
2 parents 69dbb4f + 3202d3d commit d56ae72

4 files changed

Lines changed: 270 additions & 38 deletions

File tree

‎README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,8 +22,8 @@ Supported operators:
2222
- [x] `hcat`
2323
- [x] `vcat`
2424
- [x] `norm`
25-
- [ ] Matrix-Vector product
26-
- [ ] Matrix-Matrix product
25+
- [x] Matrix-Vector product
26+
- [x] Matrix-Matrix product
2727
- [ ] Broadcasting scalar operator
2828

2929
Supported levels of AD:

‎src/reverse_mode.jl‎

Lines changed: 83 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -158,36 +158,62 @@ function _forward_eval(
158158
@j f.forward_storage[k] = tmp_sub
159159
end
160160
elseif node.index == 3 # :*
161-
tmp_prod = one(T)
162-
for c_idx in children_indices
163-
@inbounds tmp_prod *= f.forward_storage[children_arr[c_idx]]
164-
end
165-
if tmp_prod == zero(T) || N <= 2
166-
# This is inefficient if there are a lot of children.
167-
# 2 is chosen as a limit because (x*y)/y does not always
168-
# equal x for floating-point numbers. This can produce
169-
# unexpected error in partials. There's still an error when
170-
# multiplying three or more terms, but users are less likely
171-
# to complain about it.
172-
for c_idx in children_indices
173-
prod_others = one(T)
174-
for c_idx2 in children_indices
175-
(c_idx == c_idx2) && continue
176-
ix = children_arr[c_idx2]
177-
prod_others *= f.forward_storage[ix]
178-
end
179-
f.partials_storage[children_arr[c_idx]] = prod_others
161+
# Node `k` is not scalar, so we do matrix multiplication
162+
if f.sizes.ndims[k] != 0
163+
@assert N == 2
164+
idx1 = first(children_indices)
165+
idx2 = last(children_indices)
166+
@inbounds ix1 = children_arr[idx1]
167+
@inbounds ix2 = children_arr[idx2]
168+
v1 = zeros(_size(f.sizes, ix1)...)
169+
v2 = zeros(_size(f.sizes, ix2)...)
170+
for j in _eachindex(f.sizes, ix1)
171+
v1[j] = @j f.forward_storage[ix1]
172+
@j f.partials_storage[ix2] = v1[j]
173+
end
174+
for j in _eachindex(f.sizes, ix2)
175+
v2[j] = @j f.forward_storage[ix2]
176+
@j f.partials_storage[ix1] = v2[j]
177+
end
178+
v_prod = v1 * v2
179+
for j in _eachindex(f.sizes, k)
180+
@j f.forward_storage[k] = v_prod[j]
180181
end
182+
# Node `k` is scalar
181183
else
182-
# Compute all-minus-one partial derivatives by dividing from
183-
# the total product.
184+
tmp_prod = one(T)
184185
for c_idx in children_indices
185-
ix = children_arr[c_idx]
186-
f.partials_storage[ix] =
187-
tmp_prod / f.forward_storage[ix]
186+
@inbounds tmp_prod *=
187+
f.forward_storage[children_arr[c_idx]]
188188
end
189+
if tmp_prod == zero(T) || N <= 2
190+
# This is inefficient if there are a lot of children.
191+
# 2 is chosen as a limit because (x*y)/y does not always
192+
# equal x for floating-point numbers. This can produce
193+
# unexpected error in partials. There's still an error when
194+
# multiplying three or more terms, but users are less likely
195+
# to complain about it.
196+
for c_idx in children_indices
197+
prod_others = one(T)
198+
for c_idx2 in children_indices
199+
(c_idx == c_idx2) && continue
200+
ix = children_arr[c_idx2]
201+
prod_others *= f.forward_storage[ix]
202+
end
203+
f.partials_storage[children_arr[c_idx]] =
204+
prod_others
205+
end
206+
else
207+
# Compute all-minus-one partial derivatives by dividing from
208+
# the total product.
209+
for c_idx in children_indices
210+
ix = children_arr[c_idx]
211+
f.partials_storage[ix] =
212+
tmp_prod / f.forward_storage[ix]
213+
end
214+
end
215+
@inbounds f.forward_storage[k] = tmp_prod
189216
end
190-
@inbounds f.forward_storage[k] = tmp_prod
191217
elseif node.index == 4 # :^
192218
@assert N == 2
193219
idx1 = first(children_indices)
@@ -430,9 +456,39 @@ function _reverse_eval(f::_SubexpressionStorage)
430456
node = f.nodes[k]
431457
children_indices = SparseArrays.nzrange(f.adj, k)
432458
if node.type == MOI.Nonlinear.NODE_CALL_MULTIVARIATE
433-
if node.index in eachindex(DEFAULT_MULTIVARIATE_OPERATORS)
459+
if node.index in
460+
eachindex(DEFAULT_MULTIVARIATE_OPERATORS)
434461
op = DEFAULT_MULTIVARIATE_OPERATORS[node.index]
435-
if op == :vect
462+
if op == :*
463+
if f.sizes.ndims[k] != 0
464+
# Node `k` is not scalar, so we do matrix multiplication
465+
idx1 = first(children_indices)
466+
idx2 = last(children_indices)
467+
ix1 = children_arr[idx1]
468+
ix2 = children_arr[idx2]
469+
v1 = zeros(_size(f.sizes, ix1)...)
470+
v2 = zeros(_size(f.sizes, ix2)...)
471+
for j in _eachindex(f.sizes, ix1)
472+
v1[j] = @j f.forward_storage[ix1]
473+
end
474+
for j in _eachindex(f.sizes, ix2)
475+
v2[j] = @j f.forward_storage[ix2]
476+
end
477+
rev_parent = zeros(_size(f.sizes, k)...)
478+
for j in _eachindex(f.sizes, k)
479+
rev_parent[j] = @j f.reverse_storage[k]
480+
end
481+
rev_v1 = rev_parent * v2'
482+
rev_v2 = v1' * rev_parent
483+
for j in _eachindex(f.sizes, ix1)
484+
@j f.reverse_storage[ix1] = rev_v1[j]
485+
end
486+
for j in _eachindex(f.sizes, ix2)
487+
@j f.reverse_storage[ix2] = rev_v2[j]
488+
end
489+
continue
490+
end
491+
elseif op == :vect
436492
@assert _eachindex(f.sizes, k) ==
437493
eachindex(children_indices)
438494
for j in eachindex(children_indices)

‎src/sizes.jl‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -235,26 +235,29 @@ function _infer_sizes(
235235
return !iszero(sizes.ndims[children_arr[i]])
236236
end
237237
if !isnothing(first_matrix)
238-
last_matrix = findfirst(children_indices) do i
239-
return !iszero(sizes.ndims[children_arr[i]])
240-
end
241-
if sizes.ndims[last_matrix] == 0 ||
242-
sizes.ndims[first_matrix] == 0
238+
if sizes.ndims[children_arr[first(children_indices)]] == 0
243239
_add_size!(sizes, k, (1, 1))
244240
continue
245241
else
246242
_add_size!(
247243
sizes,
248244
k,
249245
(
250-
_size(sizes, first_matrix, 1),
251246
_size(
252247
sizes,
253-
last_matrix,
254-
sizes.ndims[last_matrix],
248+
children_arr[first(children_indices)],
249+
1,
250+
),
251+
_size(
252+
sizes,
253+
children_arr[last(children_indices)],
254+
sizes.ndims[children_arr[last(
255+
children_indices,
256+
)],],
255257
),
256258
),
257259
)
260+
continue
258261
end
259262
end
260263
elseif op == :^ || op == :/

‎test/ArrayDiff.jl‎

Lines changed: 174 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ function test_objective_dot_univariate()
2626
x = MOI.VariableIndex(1)
2727
ArrayDiff.set_objective(model, :(dot([$x], [$x])))
2828
evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x])
29-
MOI.initialize(evaluator, [:Grad])
29+
MOI.initialize(evaluator, [:Grad, :Hess])
3030
sizes = evaluator.backend.objective.expr.sizes
3131
@test sizes.ndims == [0, 1, 0, 1, 0]
3232
@test sizes.size_offset == [0, 1, 0, 0, 0]
@@ -356,6 +356,179 @@ function test_objective_norm_of_matrix_with_sum()
356356
return
357357
end
358358

359+
function test_objective_norm_of_product_of_matrices()
360+
model = Nonlinear.Model()
361+
x1 = MOI.VariableIndex(1)
362+
x2 = MOI.VariableIndex(2)
363+
x3 = MOI.VariableIndex(3)
364+
x4 = MOI.VariableIndex(4)
365+
Nonlinear.set_objective(model, :(norm([$x1 $x2; $x3 $x4] * [1 0; 0 1])))
366+
evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x1, x2, x3, x4])
367+
MOI.initialize(evaluator, [:Grad])
368+
sizes = evaluator.backend.objective.expr.sizes
369+
@test sizes.ndims == [0, 2, 2, 2, 0, 0, 2, 0, 0, 2, 2, 0, 0, 2, 0, 0]
370+
@test sizes.size_offset ==
371+
[0, 12, 10, 8, 0, 0, 6, 0, 0, 4, 2, 0, 0, 0, 0, 0]
372+
@test sizes.size == [1, 2, 1, 2, 2, 2, 1, 2, 1, 2, 2, 2, 2, 2]
373+
@test sizes.storage_offset ==
374+
[0, 1, 5, 9, 11, 12, 13, 15, 16, 17, 21, 23, 24, 25, 27, 28, 29]
375+
x1 = 1.0
376+
x2 = 2.0
377+
x3 = 3.0
378+
x4 = 4.0
379+
@test MOI.eval_objective(evaluator, [x1, x2, x3, x4]) == sqrt(30.0)
380+
g = ones(4)
381+
MOI.eval_objective_gradient(evaluator, g, [x1, x2, x3, x4])
382+
@test g == [
383+
1.0 / sqrt(30.0),
384+
2.0 / sqrt(30.0),
385+
3.0 / sqrt(30.0),
386+
4.0 / sqrt(30.0),
387+
]
388+
return
389+
end
390+
391+
function test_objective_norm_of_product_of_matrices_with_sum()
392+
model = Nonlinear.Model()
393+
x1 = MOI.VariableIndex(1)
394+
x2 = MOI.VariableIndex(2)
395+
x3 = MOI.VariableIndex(3)
396+
x4 = MOI.VariableIndex(4)
397+
Nonlinear.set_objective(
398+
model,
399+
:(norm(([$x1 $x2; $x3 $x4] + [1 1; 1 1]) * [1 0; 0 1])),
400+
)
401+
evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x1, x2, x3, x4])
402+
MOI.initialize(evaluator, [:Grad])
403+
sizes = evaluator.backend.objective.expr.sizes
404+
@test sizes.ndims == [
405+
0,
406+
2,
407+
2,
408+
2,
409+
2,
410+
0,
411+
0,
412+
2,
413+
0,
414+
0,
415+
2,
416+
2,
417+
0,
418+
0,
419+
2,
420+
0,
421+
0,
422+
2,
423+
2,
424+
0,
425+
0,
426+
2,
427+
0,
428+
0,
429+
]
430+
@test sizes.size_offset == [
431+
0,
432+
20,
433+
18,
434+
16,
435+
14,
436+
0,
437+
0,
438+
12,
439+
0,
440+
0,
441+
10,
442+
8,
443+
0,
444+
0,
445+
6,
446+
0,
447+
0,
448+
4,
449+
2,
450+
0,
451+
0,
452+
0,
453+
0,
454+
0,
455+
]
456+
@test sizes.size ==
457+
[1, 2, 1, 2, 2, 2, 1, 2, 1, 2, 2, 2, 1, 2, 1, 2, 2, 2, 2, 2, 2, 2]
458+
@test sizes.storage_offset == [
459+
0,
460+
1,
461+
5,
462+
9,
463+
13,
464+
15,
465+
16,
466+
17,
467+
19,
468+
20,
469+
21,
470+
25,
471+
27,
472+
28,
473+
29,
474+
31,
475+
32,
476+
33,
477+
37,
478+
39,
479+
40,
480+
41,
481+
43,
482+
44,
483+
45,
484+
]
485+
x1 = 1.0
486+
x2 = 2.0
487+
x3 = 3.0
488+
x4 = 4.0
489+
@test MOI.eval_objective(evaluator, [x1, x2, x3, x4]) == sqrt(54.0)
490+
g = ones(4)
491+
MOI.eval_objective_gradient(evaluator, g, [x1, x2, x3, x4])
492+
@test g == [
493+
2.0 / sqrt(54.0),
494+
3.0 / sqrt(54.0),
495+
4.0 / sqrt(54.0),
496+
5.0 / sqrt(54.0),
497+
]
498+
return
499+
end
500+
501+
function test_objective_norm_of_mtx_vector_product()
502+
model = Nonlinear.Model()
503+
x1 = MOI.VariableIndex(1)
504+
x2 = MOI.VariableIndex(2)
505+
x3 = MOI.VariableIndex(3)
506+
x4 = MOI.VariableIndex(4)
507+
Nonlinear.set_objective(model, :(norm(([$x1 $x2; $x3 $x4] * [1; 1]))))
508+
evaluator = Nonlinear.Evaluator(model, ArrayDiff.Mode(), [x1, x2, x3, x4])
509+
MOI.initialize(evaluator, [:Grad])
510+
sizes = evaluator.backend.objective.expr.sizes
511+
@test sizes.ndims == [0, 2, 2, 2, 0, 0, 2, 0, 0, 2, 0, 0]
512+
@test sizes.size_offset == [0, 8, 6, 4, 0, 0, 2, 0, 0, 0, 0, 0]
513+
@test sizes.size == [2, 1, 1, 2, 1, 2, 2, 2, 2, 1]
514+
@test sizes.storage_offset ==
515+
[0, 1, 3, 7, 9, 10, 11, 13, 14, 15, 17, 18, 19]
516+
x1 = 1.0
517+
x2 = 2.0
518+
x3 = 3.0
519+
x4 = 4.0
520+
@test MOI.eval_objective(evaluator, [x1, x2, x3, x4]) == sqrt(58.0)
521+
g = ones(4)
522+
MOI.eval_objective_gradient(evaluator, g, [x1, x2, x3, x4])
523+
@test g == [
524+
3.0 / sqrt(58.0),
525+
3.0 / sqrt(58.0),
526+
7.0 / sqrt(58.0),
527+
7.0 / sqrt(58.0),
528+
]
529+
return
530+
end
531+
359532
end # module
360533

361534
TestArrayDiff.runtests()

0 commit comments

Comments
 (0)