182 lines
2.5 KiB
Python
182 lines
2.5 KiB
Python
"""isort:skip_file"""
|
|
# Import order is significant here.
|
|
|
|
from ..impl import (
|
|
ir,
|
|
builtin,
|
|
)
|
|
from . import libdevice
|
|
from .core import (
|
|
abs,
|
|
arange,
|
|
argmin,
|
|
argmax,
|
|
atomic_add,
|
|
atomic_and,
|
|
atomic_cas,
|
|
atomic_max,
|
|
atomic_min,
|
|
atomic_or,
|
|
atomic_xchg,
|
|
atomic_xor,
|
|
bfloat16,
|
|
block_type,
|
|
broadcast,
|
|
broadcast_to,
|
|
cat,
|
|
cdiv,
|
|
constexpr,
|
|
cos,
|
|
debug_barrier,
|
|
dot,
|
|
dtype,
|
|
exp,
|
|
fdiv,
|
|
float16,
|
|
float32,
|
|
float64,
|
|
float8,
|
|
function_type,
|
|
int1,
|
|
int16,
|
|
int32,
|
|
int64,
|
|
int8,
|
|
load,
|
|
log,
|
|
max,
|
|
max_contiguous,
|
|
maximum,
|
|
min,
|
|
minimum,
|
|
multiple_of,
|
|
num_programs,
|
|
pi32_t,
|
|
pointer_type,
|
|
printf,
|
|
program_id,
|
|
ravel,
|
|
reshape,
|
|
sigmoid,
|
|
sin,
|
|
softmax,
|
|
sqrt,
|
|
store,
|
|
sum,
|
|
swizzle2d,
|
|
tensor,
|
|
trans,
|
|
triton,
|
|
uint16,
|
|
uint32,
|
|
uint64,
|
|
uint8,
|
|
umulhi,
|
|
view,
|
|
void,
|
|
where,
|
|
xor_sum,
|
|
zeros,
|
|
zeros_like,
|
|
)
|
|
from .random import (
|
|
pair_uniform_to_normal,
|
|
philox,
|
|
philox_impl,
|
|
rand,
|
|
rand4x,
|
|
randint,
|
|
randint4x,
|
|
randn,
|
|
randn4x,
|
|
uint32_to_uniform_float,
|
|
)
|
|
|
|
|
|
__all__ = [
|
|
"abs",
|
|
"arange",
|
|
"argmin",
|
|
"argmax",
|
|
"atomic_add",
|
|
"atomic_and",
|
|
"atomic_cas",
|
|
"atomic_max",
|
|
"atomic_min",
|
|
"atomic_or",
|
|
"atomic_xchg",
|
|
"atomic_xor",
|
|
"bfloat16",
|
|
"block_type",
|
|
"broadcast",
|
|
"broadcast_to",
|
|
"builtin",
|
|
"cat",
|
|
"cdiv",
|
|
"constexpr",
|
|
"cos",
|
|
"debug_barrier",
|
|
"dot",
|
|
"dtype",
|
|
"exp",
|
|
"fdiv",
|
|
"float16",
|
|
"float32",
|
|
"float64",
|
|
"float8",
|
|
"function_type",
|
|
"int1",
|
|
"int16",
|
|
"int32",
|
|
"int64",
|
|
"int8",
|
|
"ir",
|
|
"libdevice",
|
|
"load",
|
|
"log",
|
|
"max",
|
|
"max_contiguous",
|
|
"maximum",
|
|
"min",
|
|
"minimum",
|
|
"multiple_of",
|
|
"num_programs",
|
|
"pair_uniform_to_normal",
|
|
"philox",
|
|
"philox_impl",
|
|
"pi32_t",
|
|
"pointer_type",
|
|
"printf",
|
|
"program_id",
|
|
"rand",
|
|
"rand4x",
|
|
"randint",
|
|
"randint4x",
|
|
"randn",
|
|
"randn4x",
|
|
"ravel",
|
|
"reshape",
|
|
"sigmoid",
|
|
"sin",
|
|
"softmax",
|
|
"sqrt",
|
|
"store",
|
|
"sum",
|
|
"swizzle2d",
|
|
"tensor",
|
|
"trans",
|
|
"triton",
|
|
"uint16",
|
|
"uint32",
|
|
"uint32_to_uniform_float",
|
|
"uint64",
|
|
"uint8",
|
|
"umulhi",
|
|
"view",
|
|
"void",
|
|
"where",
|
|
"xor_sum",
|
|
"zeros",
|
|
"zeros_like",
|
|
]
|