pytorch

1
## @package algebra
2
# Module caffe2.python.helpers.algebra
3

4

5

6

7

8

9
def transpose(model, blob_in, blob_out, use_cudnn=False, **kwargs):
10
    """Transpose."""
11
    if use_cudnn:
12
        kwargs['engine'] = 'CUDNN'
13
    return model.net.Transpose(blob_in, blob_out, **kwargs)
14

15

16
def sum(model, blob_in, blob_out, **kwargs):
17
    """Sum"""
18
    return model.net.Sum(blob_in, blob_out, **kwargs)
19

20

21
def reduce_sum(model, blob_in, blob_out, **kwargs):
22
    """ReduceSum"""
23
    return model.net.ReduceSum(blob_in, blob_out, **kwargs)
24

25

26
def sub(model, blob_in, blob_out, **kwargs):
27
    """Subtract"""
28
    return model.net.Sub(blob_in, blob_out, **kwargs)
29

30

31
def mat_mul(model, blob_in, blob_out, **kwargs):
32
    """Matrix multiplication"""
33
    return model.net.MatMul(blob_in, blob_out, **kwargs)
34

35

36
def arg_min(model, blob_in, blob_out, **kwargs):
37
    """ArgMin"""
38
    return model.net.ArgMin(blob_in, blob_out, **kwargs)
39

40
def batch_mat_mul(model, blob_in, blob_out,
41
                  enable_tensor_core=False, **kwargs):
42
    if enable_tensor_core:
43
        kwargs['engine'] = 'TENSORCORE'
44

45
    return model.net.BatchMatMul(blob_in, blob_out, **kwargs)
46

47
def sparse_lengths_sum_4bit_rowwise_sparse(model, blob_in, blob_out, **kwargs):
48
    return model.net.SparseLengthsSum4BitRowwiseSparse(blob_in, blob_out, **kwargs)
49