File: find_op_test.py

package info (click to toggle)
pytorch 1.13.1%2Bdfsg-4
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 139,252 kB
  • sloc: cpp: 1,100,274; python: 706,454; ansic: 83,052; asm: 7,618; java: 3,273; sh: 2,841; javascript: 612; makefile: 323; xml: 269; ruby: 185; yacc: 144; objc: 68; lex: 44
file content (51 lines) | stat: -rw-r--r-- 1,316 bytes parent folder | download | duplicates (2)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51





from caffe2.python import core
import caffe2.python.hypothesis_test_util as hu
import caffe2.python.serialized_test.serialized_test_util as serial

import hypothesis.strategies as st
from hypothesis import given, settings
import numpy as np


class TestFindOperator(serial.SerializedTestCase):

    @given(n=st.sampled_from([1, 4, 8, 31, 79, 150]),
        idxsize=st.sampled_from([2, 4, 8, 1000, 5000]),
        **hu.gcs)
    @settings(deadline=10000)
    def test_find(self, n, idxsize, gc, dc):
        maxval = 10

        def findop(idx, X):
            res = []
            for j in list(X.flatten()):
                i = np.where(idx == j)[0]
                if len(i) == 0:
                    res.append(-1)
                else:
                    res.append(i[-1])

            print("Idx: {} X: {}".format(idx, X))
            print("Res: {}".format(res))
            return [np.array(res).astype(np.int32)]

        X = (np.random.rand(n) * maxval).astype(np.int32)
        idx = (np.random.rand(idxsize) * maxval).astype(np.int32)

        op = core.CreateOperator(
            "Find",
            ["idx", "X"],
            ["y"],
        )

        self.assertReferenceChecks(
            device_option=gc,
            op=op,
            inputs=[idx, X],
            reference=findop,
        )