op_shape_test.cpp 46.9 KB
Newer Older
Paul's avatar
Paul committed
1
2
3
4
#include <migraphx/program.hpp>
#include <migraphx/iterator_for.hpp>
#include <migraphx/instruction.hpp>
#include <migraphx/operators.hpp>
5
6
7
#include <sstream>
#include "test.hpp"

Paul's avatar
Paul committed
8
template <class... Ts>
Paul's avatar
Paul committed
9
void expect_shape(const migraphx::shape& expected, const migraphx::operation& op, Ts... xs)
10
{
Paul's avatar
Paul committed
11
12
13
    migraphx::program p;
    std::vector<migraphx::shape> shapes{xs...};
    std::vector<migraphx::instruction_ref> args(shapes.size());
Paul's avatar
Paul committed
14
15
    std::transform(
        shapes.begin(), shapes.end(), args.begin(), [&](auto&& s) { return p.add_outline(s); });
16
    p.add_instruction(op, args);
Paul's avatar
Paul committed
17
18
    if(p.get_shape() != expected)
    {
19
20
        std::cout << "FAILED: Incorrect shape for " << op.name() << ": ";
        std::cout << expected << " != " << p.get_shape() << std::endl;
Paul's avatar
Paul committed
21
        for(auto&& s : shapes)
22
23
24
25
            std::cout << "    " << s << std::endl;
    }
}

Paul's avatar
Paul committed
26
template <class... Ts>
Paul's avatar
Paul committed
27
void throws_shape(const migraphx::operation& op, Ts... xs)
28
{
Paul's avatar
Paul committed
29
30
31
    migraphx::program p;
    std::vector<migraphx::shape> shapes{xs...};
    std::vector<migraphx::instruction_ref> args(shapes.size());
Paul's avatar
Paul committed
32
33
    std::transform(
        shapes.begin(), shapes.end(), args.begin(), [&](auto&& s) { return p.add_outline(s); });
34
    bool thrown = test::throws([&] { p.add_instruction(op, args); });
Paul's avatar
Paul committed
35
36
    if(not thrown)
    {
37
        std::cout << "FAILED: No error found for " << op.name() << ": ";
Paul's avatar
Paul committed
38
        for(auto&& s : shapes)
39
40
41
42
            std::cout << "    " << s << std::endl;
    }
}

Paul's avatar
Paul committed
43
44
45
46
template <class...>
struct always_false : std::false_type
{
};
47

Paul's avatar
Paul committed
48
template <class... Ts>
Paul's avatar
Paul committed
49
void throws_shape(const migraphx::shape&, Ts...)
50
{
Paul's avatar
Paul committed
51
52
    static_assert(always_false<Ts...>{},
                  "An expected shape should not be passed to throws_shape function");
53
54
}

Paul's avatar
Paul committed
55
TEST_CASE(batch_norm_inference_shape)
56
57
{
    const size_t channels = 3;
Paul's avatar
Paul committed
58
59
60
61
62
    migraphx::shape s{migraphx::shape::float_type, {4, channels, 3, 3}};
    migraphx::shape vars{migraphx::shape::float_type, {channels}};
    expect_shape(s, migraphx::op::batch_norm_inference{}, s, vars, vars, vars, vars);
    throws_shape(migraphx::op::batch_norm_inference{}, s);
    throws_shape(migraphx::op::batch_norm_inference{}, s, vars, vars, vars, vars, vars);
63
64
}

Paul's avatar
Paul committed
65
TEST_CASE(convolution_shape)
66
{
Paul's avatar
Paul committed
67
68
69
70
71
72
73
74
75
76
    migraphx::shape output{migraphx::shape::float_type, {4, 4, 1, 1}};
    migraphx::shape input{migraphx::shape::float_type, {4, 3, 3, 3}};
    migraphx::shape weights{migraphx::shape::float_type, {4, 3, 3, 3}};
    expect_shape(output, migraphx::op::convolution{}, input, weights);
    throws_shape(migraphx::op::convolution{}, input);

    migraphx::shape input2{migraphx::shape::float_type, {3, 3}};
    migraphx::shape weights2{migraphx::shape::float_type, {3, 3}};
    throws_shape(migraphx::op::convolution{}, input2, weights2);
    throws_shape(migraphx::op::convolution{}, input2, weights);
77
78
}

Paul's avatar
Paul committed
79
TEST_CASE(transpose_shape)
80
{
Paul's avatar
Paul committed
81
82
83
84
85
    migraphx::shape input{migraphx::shape::float_type, {2, 2}};
    migraphx::shape output{migraphx::shape::float_type, {2, 2}, {1, 2}};
    expect_shape(input, migraphx::op::transpose{{0, 1}}, input);
    expect_shape(output, migraphx::op::transpose{{1, 0}}, input);
    throws_shape(migraphx::op::transpose{{1, 2}}, input);
86
87
}

Paul's avatar
Paul committed
88
TEST_CASE(contiguous_shape)
89
{
Paul's avatar
Paul committed
90
91
92
93
    migraphx::shape output{migraphx::shape::float_type, {2, 2}};
    migraphx::shape input{migraphx::shape::float_type, {2, 2}, {1, 2}};
    expect_shape(output, migraphx::op::contiguous{}, input);
    throws_shape(migraphx::op::contiguous{}, input, input);
Paul's avatar
Paul committed
94

Paul's avatar
Paul committed
95
96
    migraphx::shape single{migraphx::shape::float_type, {2}};
    expect_shape(single, migraphx::op::contiguous{}, single);
97
98
}

Paul's avatar
Paul committed
99
TEST_CASE(reshape_shape)
100
{
Paul's avatar
Paul committed
101
    migraphx::shape input{migraphx::shape::float_type, {24, 1, 1, 1}};
Paul's avatar
Paul committed
102
103
    for(auto&& new_shape :
        std::vector<std::vector<int64_t>>{{8, 3, 1, 1}, {1, 3, 4, 2}, {1, 3, 4, 2}})
104
105
106
    {
        std::vector<std::size_t> lens(new_shape.size());
        std::copy(new_shape.begin(), new_shape.end(), lens.begin());
Paul's avatar
Paul committed
107
108
        migraphx::shape output{migraphx::shape::float_type, lens};
        expect_shape(output, migraphx::op::reshape{new_shape}, input);
109
110
    }

Shucai Xiao's avatar
Shucai Xiao committed
111
112
    for(auto&& new_shape :
        std::vector<std::vector<int64_t>>{{8, 3, 2, 2}, {1, 3, -1, -1}, {3, 0, 0}, {3, 2, 0}})
113
    {
Paul's avatar
Paul committed
114
        throws_shape(migraphx::op::reshape{new_shape}, input);
115
    }
Shucai Xiao's avatar
Shucai Xiao committed
116

Shucai Xiao's avatar
Shucai Xiao committed
117
    std::vector<std::pair<std::vector<int64_t>, migraphx::shape>> minus1_tests{
Shucai Xiao's avatar
Shucai Xiao committed
118
119
120
121
122
123
124
125
        {{2, -1, 3}, {migraphx::shape::float_type, {2, 4, 3}}},
        {{0, -1, 0}, {migraphx::shape::float_type, {24, 1, 1}}},
        {{2, -1, 0}, {migraphx::shape::float_type, {2, 12, 1}}},
        {{0, 0, -1}, {migraphx::shape::float_type, {24, 1, 1}}},
        {{2, 0, -1}, {migraphx::shape::float_type, {2, 1, 12}}},
        {{-1, 2, 3}, {migraphx::shape::float_type, {4, 2, 3}}},
        {{-1, 0, 3}, {migraphx::shape::float_type, {8, 1, 3}}},
        {{-1, 0, 0}, {migraphx::shape::float_type, {24, 1, 1}}},
Shucai Xiao's avatar
Shucai Xiao committed
126
        {{-1, 3, 0}, {migraphx::shape::float_type, {8, 3, 1}}}};
Shucai Xiao's avatar
Shucai Xiao committed
127

Shucai Xiao's avatar
Shucai Xiao committed
128
    for(auto& it : minus1_tests)
Shucai Xiao's avatar
Shucai Xiao committed
129
130
131
    {
        expect_shape(it.second, migraphx::op::reshape{it.first}, input);
    }
132
133
}

Paul's avatar
Paul committed
134
TEST_CASE(flatten_shape)
135
{
Paul's avatar
Paul committed
136
137
138
    migraphx::shape input{migraphx::shape::float_type, {2, 4, 6, 8}};
    expect_shape(migraphx::shape{migraphx::shape::float_type, {1, 2 * 4 * 6 * 8}},
                 migraphx::op::flatten{0},
Scott Thornton's avatar
Scott Thornton committed
139
                 input);
Paul's avatar
Paul committed
140
141
142
143
144
145
146
147
148
    expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 4 * 6 * 8}},
                 migraphx::op::flatten{1},
                 input);
    expect_shape(migraphx::shape{migraphx::shape::float_type, {2 * 4, 6 * 8}},
                 migraphx::op::flatten{2},
                 input);
    expect_shape(migraphx::shape{migraphx::shape::float_type, {2 * 4 * 6, 8}},
                 migraphx::op::flatten{3},
                 input);
Paul's avatar
Paul committed
149
150
    expect_shape(migraphx::shape{migraphx::shape::float_type, {2 * 4 * 6 * 8, 1}},
                 migraphx::op::flatten{4},
Scott Thornton's avatar
Scott Thornton committed
151
                 input);
Paul's avatar
Paul committed
152
    throws_shape(migraphx::op::flatten{5}, input);
153
154
}

Paul's avatar
Paul committed
155
TEST_CASE(slice_shape)
Scott Thornton's avatar
Scott Thornton committed
156
{
Paul's avatar
Paul committed
157
158
159
    migraphx::shape input{migraphx::shape::int32_type, {2, 2, 3}};
    expect_shape(migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}},
                 migraphx::op::slice{{2}, {1}, {3}},
Scott Thornton's avatar
Scott Thornton committed
160
                 input);
Paul's avatar
Paul committed
161
162
    expect_shape(migraphx::shape{migraphx::shape::int32_type, {2, 2, 2}, {6, 3, 1}},
                 migraphx::op::slice{{0, 1, 2}, {0, 0, 1}, {2, 2, 3}},
Scott Thornton's avatar
Scott Thornton committed
163
                 input);
Paul's avatar
Paul committed
164
165
    expect_shape(migraphx::shape{migraphx::shape::int32_type, {2, 2, 1}, {6, 3, 1}},
                 migraphx::op::slice{{2}, {2}, {10}},
Scott Thornton's avatar
Scott Thornton committed
166
167
                 input);
}
Scott Thornton's avatar
Scott Thornton committed
168

wsttiger's avatar
wsttiger committed
169
TEST_CASE(multibroadcast)
Scott Thornton's avatar
Scott Thornton committed
170
171
{
    {
172
        std::vector<std::size_t> lens{4, 2, 5, 3};
Paul's avatar
Paul committed
173
174
175
        migraphx::shape input{migraphx::shape::float_type, {2, 1, 3}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, lens, {0, 3, 0, 1}},
                     migraphx::op::multibroadcast{lens},
176
                     input);
Scott Thornton's avatar
Scott Thornton committed
177
178
    }
    {
179
        std::vector<std::size_t> lens{4, 2, 5, 3};
Paul's avatar
Paul committed
180
181
182
        migraphx::shape input{migraphx::shape::float_type, {2, 1, 1}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, lens, {0, 1, 0, 0}},
                     migraphx::op::multibroadcast{lens},
183
                     input);
Scott Thornton's avatar
Scott Thornton committed
184
185
    }
    {
186
        std::vector<std::size_t> lens{4, 2, 5, 3};
Paul's avatar
Paul committed
187
188
189
        migraphx::shape input{migraphx::shape::float_type, {5, 1}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, lens, {0, 0, 1, 0}},
                     migraphx::op::multibroadcast{lens},
190
                     input);
Scott Thornton's avatar
Scott Thornton committed
191
192
    }
    {
193
        std::vector<std::size_t> lens{4, 2, 5, 3};
Paul's avatar
Paul committed
194
195
196
        migraphx::shape input{migraphx::shape::float_type, {4, 1, 1, 1}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, lens, {1, 0, 0, 0}},
                     migraphx::op::multibroadcast{lens},
197
198
199
200
                     input);
    }
    {
        std::vector<std::size_t> lens{4, 2, 5, 3};
Paul's avatar
Paul committed
201
202
203
        migraphx::shape input{migraphx::shape::float_type, {3}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, lens, {0, 0, 0, 1}},
                     migraphx::op::multibroadcast{lens},
204
205
                     input);
    }
206
207
    {
        std::vector<std::size_t> lens{4, 4, 1, 3};
Paul's avatar
Paul committed
208
209
210
        migraphx::shape input{migraphx::shape::float_type, {4, 1, 3}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, lens, {0, 3, 3, 1}},
                     migraphx::op::multibroadcast{lens},
211
212
                     input);
    }
213
214
    {
        std::vector<std::size_t> lens{4, 1, 1, 3};
Paul's avatar
Paul committed
215
216
217
        migraphx::shape input{migraphx::shape::float_type, {4, 1, 1, 1}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, lens, {1, 1, 1, 0}},
                     migraphx::op::multibroadcast{lens},
218
219
220
221
                     input);
    }
    {
        std::vector<std::size_t> lens{4, 1, 3};
Paul's avatar
Paul committed
222
223
        migraphx::shape input{migraphx::shape::float_type, {4, 1, 1, 1}};
        throws_shape(migraphx::op::multibroadcast{lens}, input);
Scott Thornton's avatar
Scott Thornton committed
224
225
    }
    {
226
        std::vector<std::size_t> lens{4, 1, 3};
Paul's avatar
Paul committed
227
228
        migraphx::shape input{migraphx::shape::float_type, {}};
        throws_shape(migraphx::op::multibroadcast{lens}, input);
Scott Thornton's avatar
Scott Thornton committed
229
230
231
    }
}

Shucai Xiao's avatar
Shucai Xiao committed
232
TEST_CASE(gather)
233
234
235
236
{
    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        migraphx::shape indices{migraphx::shape::int32_type, {2, 3}};
237
        int axis = 1;
238
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 2, 3, 4, 5}},
Shucai Xiao's avatar
Shucai Xiao committed
239
240
241
                     migraphx::op::gather{axis},
                     input,
                     indices);
242
243
    }

244
245
246
247
    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        migraphx::shape indices{migraphx::shape::int32_type, {2, 3}};
        int axis = -4;
248
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 3, 4, 5}},
249
250
251
252
253
                     migraphx::op::gather{axis},
                     input,
                     indices);
    }

254
255
256
257
258
259
260
261
262
263
264
265
    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        migraphx::shape indices{migraphx::shape::int32_type, {1}};
        int axis = -4;
        expect_shape(migraphx::shape{migraphx::shape::float_type, {1, 3, 4, 5}},
                     migraphx::op::gather{axis},
                     input,
                     indices);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
266
        migraphx::shape indices{migraphx::shape::int32_type};
267
268
269
270
271
272
273
274
275
        int axis = -4;
        expect_shape(migraphx::shape{migraphx::shape::float_type, {3, 4, 5}},
                     migraphx::op::gather{axis},
                     input,
                     indices);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
276
        migraphx::shape indices{migraphx::shape::int32_type};
277
278
279
280
281
282
283
284
285
        int axis = 3;
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 4}},
                     migraphx::op::gather{axis},
                     input,
                     indices);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {3}};
286
        migraphx::shape indices{migraphx::shape::int32_type};
287
        int axis = 0;
288
        expect_shape(migraphx::shape{migraphx::shape::float_type},
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
                     migraphx::op::gather{axis},
                     input,
                     indices);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {3}};
        migraphx::shape indices{migraphx::shape::int32_type, {1}};
        int axis = 0;
        expect_shape(migraphx::shape{migraphx::shape::float_type, {1}},
                     migraphx::op::gather{axis},
                     input,
                     indices);
    }

304
305
306
    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        migraphx::shape indices{migraphx::shape::int32_type, {2, 3}};
307
        int axis = 4;
308
309
        throws_shape(migraphx::op::gather{axis}, input, indices);
    }
310
311
312
313
314
315
316

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        migraphx::shape indices{migraphx::shape::int32_type, {2, 3}};
        int axis = -5;
        throws_shape(migraphx::op::gather{axis}, input, indices);
    }
317
318
}

319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
TEST_CASE(logsoftmax)
{
    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        int axis = 0;
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 4, 5}},
                     migraphx::op::logsoftmax{axis},
                     input);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        int axis = 1;
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 4, 5}},
                     migraphx::op::logsoftmax{axis},
                     input);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        int axis = 2;
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 4, 5}},
                     migraphx::op::logsoftmax{axis},
                     input);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        int axis = 3;
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 4, 5}},
                     migraphx::op::logsoftmax{axis},
                     input);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        int axis = 4;
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 4, 5}},
                     migraphx::op::logsoftmax{axis},
                     input);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        int axis = 5;
        throws_shape(migraphx::op::logsoftmax{axis}, input);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        int axis = -1;
        throws_shape(migraphx::op::logsoftmax{axis}, input);
    }
}

374
375
// 2 inputs arguments
TEST_CASE(matmul)
376
{
377
378
379
380
381
382
383
384
    {
        migraphx::shape s_m1{migraphx::shape::float_type, {5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {5}};
385
386
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 2}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2);
387
388
389
    }

    {
390
391
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5}};
392
393
394
395
        throws_shape(migraphx::op::dot{}, s_m1, s_m2);
    }

    {
396
397
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 4}};
398
        expect_shape(
399
            migraphx::shape{migraphx::shape::float_type, {1, 4}}, migraphx::op::dot{}, s_m1, s_m2);
400
401
402
    }

    {
403
404
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {4, 4}};
405
406
407
408
        throws_shape(migraphx::op::dot{}, s_m1, s_m2);
    }

    {
409
410
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {6, 5, 4}};
411
412
413
414
        throws_shape(migraphx::op::dot{}, s_m1, s_m2);
    }

    {
415
416
        migraphx::shape s_m1{migraphx::shape::float_type, {6, 1, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {6, 5, 4}};
Shucai Xiao's avatar
Shucai Xiao committed
417
418
419
420
        expect_shape(migraphx::shape{migraphx::shape::float_type, {6, 1, 4}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2);
421
422
423
    }

    {
424
425
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 6, 1, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 6, 5, 4}};
Shucai Xiao's avatar
Shucai Xiao committed
426
427
428
429
        expect_shape(migraphx::shape{migraphx::shape::float_type, {1, 6, 1, 4}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2);
430
431
    }

432
433
434
    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
Shucai Xiao's avatar
Shucai Xiao committed
435
436
        expect_shape(
            migraphx::shape{migraphx::shape::float_type, {4, 8}}, migraphx::op::dot{}, s_m1, s_m2);
437
438
439
440
441
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 1}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 1}};
Shucai Xiao's avatar
Shucai Xiao committed
442
443
        expect_shape(
            migraphx::shape{migraphx::shape::float_type, {1, 1}}, migraphx::op::dot{}, s_m1, s_m2);
444
445
446
447
448
449
450
451
452
453
454
455
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 7}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, {1, 4, 7}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2);
    }

    {
456
457
458
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 1, 5, 7}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2);
459
460
461
462
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 1, 4, 5}};
463
464
465
466
467
468
469
470
471
472
473
474
475
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 2, 5, 7}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2);
    }
}

// 3 input arguments
TEST_CASE(gemm)
{
    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
476
477
478
    }

    {
479
480
481
482
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1, 1}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
483
484
485
486
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
487
488
489
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {8}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
490
491
    }

492
    {
493
494
495
496
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {4, 1}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
497
498
499
    }

    {
500
501
502
503
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 6}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {4, 8}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
504
505
    }

506
    {
507
508
509
510
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {4}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
511
512
513
514
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
515
516
517
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {4, 8}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, {4, 8}},
518
519
                     migraphx::op::dot{},
                     s_m1,
520
521
                     s_m2,
                     s_m3);
522
523
    }

524
    {
525
526
527
528
529
530
531
532
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1, 4, 8}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, {1, 4, 8}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2,
                     s_m3);
533
534
535
    }

    {
536
537
538
539
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 4, 6}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1, 4, 8}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
540
541
542
    }

    {
543
544
545
546
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {4, 8}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
547
548
549
    }

    {
550
551
552
553
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
554
555
556
    }
}

557
558
559
560
561
562
563
// 3 input arguments
TEST_CASE(gemm)
{
    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type};
Shucai Xiao's avatar
Shucai Xiao committed
564
565
566
567
568
        expect_shape(migraphx::shape{migraphx::shape::float_type, {4, 8}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2,
                     s_m3);
569
570
571
572
573
574
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1}};
Shucai Xiao's avatar
Shucai Xiao committed
575
576
577
578
579
        expect_shape(migraphx::shape{migraphx::shape::float_type, {4, 8}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2,
                     s_m3);
580
581
582
583
584
585
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1, 1}};
Shucai Xiao's avatar
Shucai Xiao committed
586
587
588
589
590
        expect_shape(migraphx::shape{migraphx::shape::float_type, {4, 8}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2,
                     s_m3);
591
592
593
594
595
596
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {8}};
Shucai Xiao's avatar
Shucai Xiao committed
597
598
599
600
601
        expect_shape(migraphx::shape{migraphx::shape::float_type, {4, 8}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2,
                     s_m3);
602
603
604
605
606
607
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {4, 1}};
Shucai Xiao's avatar
Shucai Xiao committed
608
609
610
611
612
        expect_shape(migraphx::shape{migraphx::shape::float_type, {4, 8}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2,
                     s_m3);
613
614
615
616
617
618
619
620
621
622
623
624
625
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {4}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {4, 8}};
Shucai Xiao's avatar
Shucai Xiao committed
626
627
628
629
630
        expect_shape(migraphx::shape{migraphx::shape::float_type, {4, 8}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2,
                     s_m3);
631
632
633
634
635
636
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1, 8}};
Shucai Xiao's avatar
Shucai Xiao committed
637
638
639
640
641
        expect_shape(migraphx::shape{migraphx::shape::float_type, {4, 8}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2,
                     s_m3);
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {4, 7}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {3, 8}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {3, 7}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1, 4, 8}};
Shucai Xiao's avatar
Shucai Xiao committed
669
670
671
672
673
        expect_shape(migraphx::shape{migraphx::shape::float_type, {1, 4, 8}},
                     migraphx::op::dot{},
                     s_m1,
                     s_m2,
                     s_m3);
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {4, 8}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1, 8}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {1, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1, 4, 1}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {2, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1, 4, 8}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {2, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {1, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type, {1, 4, 8}};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }

    {
        migraphx::shape s_m1{migraphx::shape::float_type, {2, 4, 5}};
        migraphx::shape s_m2{migraphx::shape::float_type, {2, 5, 8}};
        migraphx::shape s_m3{migraphx::shape::float_type};
        throws_shape(migraphx::op::dot{}, s_m1, s_m2, s_m3);
    }
}

726
727
728
729
730
731
732
733
TEST_CASE(rnn)
{
    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
Shucai Xiao's avatar
Shucai Xiao committed
734
        float clip              = 0.0f;
735
736
737
738
739
740
741

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
742
743
744
745
746
747
748
749
750
751
        expect_shape(
            migraphx::shape{migraphx::shape::float_type,
                            {seq_len, num_dirct, batch_size, hidden_size}},
            migraphx::op::rnn{
                hidden_size, {migraphx::op::tanh{}}, migraphx::op::rnn_direction::forward, clip},
            in_shape,
            w_shape,
            r_shape,
            b_shape,
            ih_shape);
752
753
754
755
756
757
758
759
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
Shucai Xiao's avatar
Shucai Xiao committed
760
        float clip              = 0.0f;
761
762
763
764
765
766
767

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
768
769
770
771
772
773
774
775
776
777
        expect_shape(
            migraphx::shape{migraphx::shape::float_type,
                            {seq_len, num_dirct, batch_size, hidden_size}},
            migraphx::op::rnn{
                hidden_size, {migraphx::op::tanh{}}, migraphx::op::rnn_direction::reverse, clip},
            in_shape,
            w_shape,
            r_shape,
            b_shape,
            ih_shape);
778
779
780
781
782
783
784
785
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 2;
Shucai Xiao's avatar
Shucai Xiao committed
786
        float clip              = 0.0f;
787
788
789
790
791
792
793

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
794
795
796
797
798
799
800
801
802
803
804
        expect_shape(migraphx::shape{migraphx::shape::float_type,
                                     {seq_len, num_dirct, batch_size, hidden_size}},
                     migraphx::op::rnn{hidden_size,
                                       {migraphx::op::tanh{}},
                                       migraphx::op::rnn_direction::bidirectional,
                                       clip},
                     in_shape,
                     w_shape,
                     r_shape,
                     b_shape,
                     ih_shape);
805
806
807
808
809
810
811
812
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
Shucai Xiao's avatar
Shucai Xiao committed
813
        float clip              = 0.0f;
814
815
816
817
818
819
820

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
821
822
823
824
825
826
827
828
829
        throws_shape(migraphx::op::rnn{hidden_size + 1,
                                       {migraphx::op::tanh{}},
                                       migraphx::op::rnn_direction::forward,
                                       clip},
                     in_shape,
                     w_shape,
                     r_shape,
                     b_shape,
                     ih_shape);
830
831
832
833
834
835
836
837
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
Shucai Xiao's avatar
Shucai Xiao committed
838
        float clip              = 0.0f;
839
840
841
842
843
844
845

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
846
847
848
849
850
851
852
853
854
        throws_shape(migraphx::op::rnn{hidden_size,
                                       {migraphx::op::tanh{}},
                                       migraphx::op::rnn_direction::bidirectional,
                                       clip},
                     in_shape,
                     w_shape,
                     r_shape,
                     b_shape,
                     ih_shape);
855
856
857
858
859
860
861
862
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 2;
Shucai Xiao's avatar
Shucai Xiao committed
863
        float clip              = 0.0f;
864
865
866
867
868
869
870

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};
        migraphx::shape w_shape{migraphx::shape::float_type, {num_dirct, hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type, {num_dirct, hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 2 * hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
871
872
        throws_shape(
            migraphx::op::rnn{
873
                hidden_size, {migraphx::op::tanh{}}, migraphx::op::rnn_direction::forward, clip},
Shucai Xiao's avatar
Shucai Xiao committed
874
875
876
877
878
            in_shape,
            w_shape,
            r_shape,
            b_shape,
            ih_shape);
879
880
881
    }
}

Shucai Xiao's avatar
Shucai Xiao committed
882
883
884
885
886
887
888
889
890
891
892
TEST_CASE(gru)
{
    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
Shucai Xiao's avatar
Shucai Xiao committed
893
894
895
896
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
Shucai Xiao's avatar
Shucai Xiao committed
897
898
899
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
900
901
902
903
904
905
906
907
908
909
        expect_shape(
            migraphx::shape{migraphx::shape::float_type,
                            {seq_len, num_dirct, batch_size, hidden_size}},
            migraphx::op::gru{
                hidden_size, {migraphx::op::tanh{}}, migraphx::op::rnn_direction::forward, clip},
            in_shape,
            w_shape,
            r_shape,
            b_shape,
            ih_shape);
Shucai Xiao's avatar
Shucai Xiao committed
910
911
912
913
914
915
916
917
918
919
920
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
Shucai Xiao's avatar
Shucai Xiao committed
921
922
923
924
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
Shucai Xiao's avatar
Shucai Xiao committed
925
926
927
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
928
929
930
931
932
933
934
935
936
937
        expect_shape(
            migraphx::shape{migraphx::shape::float_type,
                            {seq_len, num_dirct, batch_size, hidden_size}},
            migraphx::op::gru{
                hidden_size, {migraphx::op::tanh{}}, migraphx::op::rnn_direction::reverse, clip},
            in_shape,
            w_shape,
            r_shape,
            b_shape,
            ih_shape);
Shucai Xiao's avatar
Shucai Xiao committed
938
939
940
941
942
943
944
945
946
947
948
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 2;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
Shucai Xiao's avatar
Shucai Xiao committed
949
950
951
952
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
Shucai Xiao's avatar
Shucai Xiao committed
953
954
955
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
956
957
958
959
960
961
962
963
964
965
966
        expect_shape(migraphx::shape{migraphx::shape::float_type,
                                     {seq_len, num_dirct, batch_size, hidden_size}},
                     migraphx::op::gru{hidden_size,
                                       {migraphx::op::tanh{}},
                                       migraphx::op::rnn_direction::bidirectional,
                                       clip},
                     in_shape,
                     w_shape,
                     r_shape,
                     b_shape,
                     ih_shape);
Shucai Xiao's avatar
Shucai Xiao committed
967
968
969
970
971
972
973
974
975
976
977
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
Shucai Xiao's avatar
Shucai Xiao committed
978
979
980
981
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
Shucai Xiao's avatar
Shucai Xiao committed
982
983
984
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
985
986
987
988
989
990
991
992
993
        throws_shape(migraphx::op::gru{hidden_size + 1,
                                       {migraphx::op::tanh{}},
                                       migraphx::op::rnn_direction::forward,
                                       clip},
                     in_shape,
                     w_shape,
                     r_shape,
                     b_shape,
                     ih_shape);
Shucai Xiao's avatar
Shucai Xiao committed
994
995
996
997
998
999
1000
1001
1002
1003
1004
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
Shucai Xiao's avatar
Shucai Xiao committed
1005
1006
1007
1008
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
Shucai Xiao's avatar
Shucai Xiao committed
1009
1010
1011
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

Shucai Xiao's avatar
Shucai Xiao committed
1012
1013
1014
1015
1016
1017
1018
1019
1020
        throws_shape(migraphx::op::gru{hidden_size,
                                       {migraphx::op::tanh{}},
                                       migraphx::op::rnn_direction::bidirectional,
                                       clip},
                     in_shape,
                     w_shape,
                     r_shape,
                     b_shape,
                     ih_shape);
Shucai Xiao's avatar
Shucai Xiao committed
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 2;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
Shucai Xiao's avatar
Shucai Xiao committed
1032
1033
1034
1035
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
Shucai Xiao's avatar
Shucai Xiao committed
1036
1037
1038
1039
1040
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

        throws_shape(
            migraphx::op::gru{
1041
                hidden_size, {migraphx::op::tanh{}}, migraphx::op::rnn_direction::forward, clip},
Shucai Xiao's avatar
Shucai Xiao committed
1042
1043
1044
1045
1046
1047
1048
1049
            in_shape,
            w_shape,
            r_shape,
            b_shape,
            ih_shape);
    }
}

1050
1051
1052
1053
1054
1055
1056
1057
1058
1059
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
1076
1077
1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
1089
1090
1091
1092
1093
1094
1095
1096
1097
1098
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108
1109
1110
1111
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
TEST_CASE(lstm)
{
    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};

        expect_shape(
            migraphx::shape{migraphx::shape::float_type,
                            {seq_len, num_dirct, batch_size, hidden_size}},
            migraphx::op::lstm{
                hidden_size, {migraphx::op::tanh{}}, migraphx::op::rnn_direction::forward, clip},
            in_shape,
            w_shape,
            r_shape);
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

        expect_shape(
            migraphx::shape{migraphx::shape::float_type,
                            {seq_len, num_dirct, batch_size, hidden_size}},
            migraphx::op::lstm{
                hidden_size, {migraphx::op::tanh{}}, migraphx::op::rnn_direction::reverse, clip},
            in_shape,
            w_shape,
            r_shape,
            b_shape,
            ih_shape);
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 2;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

        expect_shape(migraphx::shape{migraphx::shape::float_type,
                                     {seq_len, num_dirct, batch_size, hidden_size}},
                     migraphx::op::lstm{hidden_size,
Shucai Xiao's avatar
Shucai Xiao committed
1123
1124
1125
                                        {migraphx::op::tanh{}},
                                        migraphx::op::rnn_direction::bidirectional,
                                        clip},
1126
1127
1128
1129
1130
1131
1132
1133
1134
1135
1136
1137
1138
1139
1140
1141
1142
1143
1144
1145
1146
1147
1148
1149
                     in_shape,
                     w_shape,
                     r_shape,
                     b_shape,
                     ih_shape);
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

        throws_shape(migraphx::op::lstm{hidden_size + 1,
Shucai Xiao's avatar
Shucai Xiao committed
1150
1151
1152
                                        {migraphx::op::tanh{}},
                                        migraphx::op::rnn_direction::forward,
                                        clip},
1153
1154
1155
1156
1157
1158
1159
1160
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170
1171
1172
1173
1174
1175
1176
                     in_shape,
                     w_shape,
                     r_shape,
                     b_shape,
                     ih_shape);
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 1;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

        throws_shape(migraphx::op::lstm{hidden_size,
Shucai Xiao's avatar
Shucai Xiao committed
1177
1178
1179
                                        {migraphx::op::tanh{}},
                                        migraphx::op::rnn_direction::bidirectional,
                                        clip},
1180
1181
1182
1183
1184
1185
1186
1187
1188
1189
1190
1191
1192
1193
1194
1195
1196
1197
1198
1199
1200
1201
1202
1203
1204
1205
1206
1207
1208
1209
1210
1211
1212
1213
                     in_shape,
                     w_shape,
                     r_shape,
                     b_shape,
                     ih_shape);
    }

    {
        std::size_t batch_size  = 2;
        std::size_t seq_len     = 2;
        std::size_t hidden_size = 4;
        std::size_t input_size  = 3;
        std::size_t num_dirct   = 2;
        float clip              = 0.0f;

        migraphx::shape in_shape{migraphx::shape::float_type, {seq_len, batch_size, input_size}};
        migraphx::shape w_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, input_size}};
        migraphx::shape r_shape{migraphx::shape::float_type,
                                {num_dirct, 3 * hidden_size, hidden_size}};
        migraphx::shape b_shape{migraphx::shape::float_type, {num_dirct, 6 * hidden_size}};
        migraphx::shape ih_shape{migraphx::shape::float_type, {num_dirct, batch_size, hidden_size}};

        throws_shape(
            migraphx::op::lstm{
                hidden_size, {migraphx::op::tanh{}}, migraphx::op::rnn_direction::forward, clip},
            in_shape,
            w_shape,
            r_shape,
            b_shape,
            ih_shape);
    }
}

Paul's avatar
Paul committed
1214
int main(int argc, const char* argv[]) { test::run(argc, argv); }