megenginelite-sys 1.8.2

A safe megenginelite wrapper in Rust
Documentation 
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

from ..core._imperative_rt.core2 import Const
from ..jit.tracing import is_tracing

small_tensor_cache = {}


def _get_scalar_tensor_with_value(value, dtype=None, device=None):
    global small_tensor_cache
    if is_tracing():
        ret = Const(value, dtype, device, None)
    else:
        cache_key = (value, dtype, device)
        if cache_key not in small_tensor_cache:
            ret = Const(value, dtype, device, None)
            small_tensor_cache[cache_key] = ret
        else:
            ret = small_tensor_cache[cache_key]
    return ret


def get_scalar_zero(dtype=None, device=None):
    return _get_scalar_tensor_with_value(0, dtype, device)


def get_scalar_zero_point_five(dtype=None, device=None):
    return _get_scalar_tensor_with_value(0.5, dtype, device)


def get_scalar_one(dtype=None, device=None):
    return _get_scalar_tensor_with_value(1, dtype, device)


def get_scalar_two(dtype=None, device=None):
    return _get_scalar_tensor_with_value(2, dtype, device)