Skip to content

GitLab

"example/01_gemm/gemm_xdl_fp64.cpp" did not exist on "4d40b1974e18e9215067fb4b1117213e69a2923e"
Switch branch/tag
Find file Normal viewHistoryPermalink
__init__.py 1.33 KB
Newer Older
rusty1s's avatar
first python impl  
rusty1s committed
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 
from .scatter import scatter
from .utils import gen_output


def scatter_add_(output, index, input, dim=0):
    return scatter('add', output, index, input, dim)


def scatter_add(index, input, dim=0, max_index=None, fill_value=0):
    output = gen_output(index, input, dim, max_index, fill_value)
    return scatter_add_(output, index, input, dim)


def scatter_sub_(output, index, input, dim=0):
    return scatter('sub', output, index, input, dim)


def scatter_sub(index, input, dim=0, max_index=None, fill_value=0):
    output = gen_output(index, input, dim, max_index, fill_value)
    return scatter_sub_(output, index, input, dim)


def scatter_mul_(output, index, input, dim=0):
    return scatter('mul', output, index, input, dim)


def scatter_mul(index, input, dim=0, max_index=None, fill_value=1):
    output = gen_output(index, input, dim, max_index, fill_value)
    return scatter_mul_(output, index, input, dim)


def scatter_div_(output, index, input, dim=0):
    return scatter('div', output, index, input, dim)


def scatter_div(index, input, dim=0, max_index=None, fill_value=1):
    output = gen_output(index, input, dim, max_index, fill_value)
    return scatter_div_(output, index, input, dim)


__all__ = [
    'scatter_add_', 'scatter_add', 'scatter_sub_', 'scatter_sub',
    'scatter_mul_', 'scatter_mul', 'scatter_div_', 'scatter_div'
]