Skip to content

Commit 409972a

Browse files
SophieL1SophieL1
authored
Merge pull request #20 from blegat/sl/mtxmult
Add mtx-mtx and mtx-vector multiplication
2 parents c168780 + ffb5dff commit 409972a

3 files changed

Lines changed: 265 additions & 34 deletions

File tree

‎src/reverse_mode.jl‎

Lines changed: 81 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -154,36 +154,62 @@ function _forward_eval(
154154
@j f.forward_storage[k] = tmp_sub
155155
end
156156
elseif node.index == 3 # :*
157-
tmp_prod = one(T)
158-
for c_idx in children_indices
159-
@inbounds tmp_prod *= f.forward_storage[children_arr[c_idx]]
160-
end
161-
if tmp_prod == zero(T) || N <= 2
162-
# This is inefficient if there are a lot of children.
163-
# 2 is chosen as a limit because (x*y)/y does not always
164-
# equal x for floating-point numbers. This can produce
165-
# unexpected error in partials. There's still an error when
166-
# multiplying three or more terms, but users are less likely
167-
# to complain about it.
168-
for c_idx in children_indices
169-
prod_others = one(T)
170-
for c_idx2 in children_indices
171-
(c_idx == c_idx2) && continue
172-
ix = children_arr[c_idx2]
173-
prod_others *= f.forward_storage[ix]
174-
end
175-
f.partials_storage[children_arr[c_idx]] = prod_others
157+
# Node `k` is not scalar, so we do matrix multiplication
158+
if f.sizes.ndims[k] != 0
159+
@assert N == 2
160+
idx1 = first(children_indices)
161+
idx2 = last(children_indices)
162+
@inbounds ix1 = children_arr[idx1]
163+
@inbounds ix2 = children_arr[idx2]
164+
v1 = zeros(_size(f.sizes, ix1)...)
165+
v2 = zeros(_size(f.sizes, ix2)...)
166+
for j in _eachindex(f.sizes, ix1)
167+
v1[j] = @j f.forward_storage[ix1]
168+
@j f.partials_storage[ix2] = v1[j]
169+
end
170+
for j in _eachindex(f.sizes, ix2)
171+
v2[j] = @j f.forward_storage[ix2]
172+
@j f.partials_storage[ix1] = v2[j]
173+
end
174+
v_prod = v1 * v2
175+
for j in _eachindex(f.sizes, k)
176+
@j f.forward_storage[k] = v_prod[j]
176177
end
178+
# Node `k` is scalar
177179
else
178-
# Compute all-minus-one partial derivatives by dividing from
179-
# the total product.
180+
tmp_prod = one(T)
180181
for c_idx in children_indices
181-
ix = children_arr[c_idx]
182-
f.partials_storage[ix] =
183-
tmp_prod / f.forward_storage[ix]
182+
@inbounds tmp_prod *=
183+
f.forward_storage[children_arr[c_idx]]
184184
end
185+
if tmp_prod == zero(T) || N <= 2
186+
# This is inefficient if there are a lot of children.
187+
# 2 is chosen as a limit because (x*y)/y does not always
188+
# equal x for floating-point numbers. This can produce
189+
# unexpected error in partials. There's still an error when
190+
# multiplying three or more terms, but users are less likely
191+
# to complain about it.
192+
for c_idx in children_indices
193+
prod_others = one(T)
194+
for c_idx2 in children_indices
195+
(c_idx == c_idx2) && continue
196+
ix = children_arr[c_idx2]
197+
prod_others *= f.forward_storage[ix]
198+
end
199+
f.partials_storage[children_arr[c_idx]] =
200+
prod_others
201+
end
202+
else
203+
# Compute all-minus-one partial derivatives by dividing from
204+
# the total product.
205+
for c_idx in children_indices
206+
ix = children_arr[c_idx]
207+
f.partials_storage[ix] =
208+
tmp_prod / f.forward_storage[ix]
209+
end
210+
end
211+
@inbounds f.forward_storage[k] = tmp_prod
185212
end
186-
@inbounds f.forward_storage[k] = tmp_prod
187213
elseif node.index == 4 # :^
188214
@assert N == 2
189215
idx1 = first(children_indices)
@@ -429,7 +455,36 @@ function _reverse_eval(f::_SubexpressionStorage)
429455
if node.index in
430456
eachindex(MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS)
431457
op = MOI.Nonlinear.DEFAULT_MULTIVARIATE_OPERATORS[node.index]
432-
if op == :vect
458+
if op == :*
459+
if f.sizes.ndims[k] != 0
460+
# Node `k` is not scalar, so we do matrix multiplication
461+
idx1 = first(children_indices)
462+
idx2 = last(children_indices)
463+
ix1 = children_arr[idx1]
464+
ix2 = children_arr[idx2]
465+
v1 = zeros(_size(f.sizes, ix1)...)
466+
v2 = zeros(_size(f.sizes, ix2)...)
467+
for j in _eachindex(f.sizes, ix1)
468+
v1[j] = @j f.forward_storage[ix1]
469+
end
470+
for j in _eachindex(f.sizes, ix2)
471+
v2[j] = @j f.forward_storage[ix2]
472+
end
473+
rev_parent = zeros(_size(f.sizes, k)...)
474+
for j in _eachindex(f.sizes, k)
475+
rev_parent[j] = @j f.reverse_storage[k]
476+
end
477+
rev_v1 = rev_parent * v2'
478+
rev_v2 = v1' * rev_parent
479+
for j in _eachindex(f.sizes, ix1)
480+
@j f.reverse_storage[ix1] = rev_v1[j]
481+
end
482+
for j in _eachindex(f.sizes, ix2)
483+
@j f.reverse_storage[ix2] = rev_v2[j]
484+
end
485+
continue
486+
end
487+
elseif op == :vect
433488
@assert _eachindex(f.sizes, k) ==
434489
eachindex(children_indices)
435490
for j in eachindex(children_indices)

‎src/sizes.jl‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -238,26 +238,29 @@ function _infer_sizes(
238238
return !iszero(sizes.ndims[children_arr[i]])
239239
end
240240
if !isnothing(first_matrix)
241-
last_matrix = findfirst(children_indices) do i
242-
return !iszero(sizes.ndims[children_arr[i]])
243-
end
244-
if sizes.ndims[last_matrix] == 0 ||
245-
sizes.ndims[first_matrix] == 0
241+
if sizes.ndims[children_arr[first(children_indices)]] == 0
246242
_add_size!(sizes, k, (1, 1))
247243
continue
248244
else
249245
_add_size!(
250246
sizes,
251247
k,
252248
(
253-
_size(sizes, first_matrix, 1),
254249
_size(
255250
sizes,
256-
last_matrix,
257-
sizes.ndims[last_matrix],
251+
children_arr[first(children_indices)],
252+
1,
253+
),
254+
_size(
255+
sizes,
256+
children_arr[last(children_indices)],
257+
sizes.ndims[children_arr[last(
258+
children_indices,
259+
)],],
258260
),
259261
),
260262
)
263+
continue
261264
end
262265
end
263266
elseif op == :^ || op == :/

‎test/ArrayDiff.jl‎

Lines changed: 173 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -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)