ColossalAI/colossalai/elixir/ctx/__init__.py

33 lines
1.6 KiB
Python

import torch
tensor_creation_methods = dict(tensor=torch.tensor,
sparse_coo_tensor=torch.sparse_coo_tensor,
asarray=torch.asarray,
as_tensor=torch.as_tensor,
as_strided=torch.as_strided,
from_numpy=torch.from_numpy,
from_dlpack=torch.from_dlpack,
frombuffer=torch.frombuffer,
zeros=torch.zeros,
zeros_like=torch.zeros_like,
ones=torch.ones,
ones_like=torch.ones_like,
arange=torch.arange,
range=torch.range,
linspace=torch.linspace,
logspace=torch.logspace,
eye=torch.eye,
empty=torch.empty,
empty_like=torch.empty_like,
empty_strided=torch.empty_strided,
full=torch.full,
full_like=torch.full_like,
quantize_per_tensor=torch.quantize_per_tensor,
quantize_per_channel=torch.quantize_per_channel,
dequantize=torch.dequantize,
complex=torch.complex,
polar=torch.polar,
heaviside=torch.heaviside)
from .meta_ctx import MetaContext