mirror of
https://github.com/hpcaitech/ColossalAI.git
synced 2025-09-20 09:01:06 +00:00
[misc] update pre-commit and run all files (#4752)
* [misc] update pre-commit * [misc] run pre-commit * [misc] remove useless configuration files * [misc] ignore cuda for clang-format
This commit is contained in:
@@ -5,4 +5,4 @@ from ...registry import meta_patched_function
|
||||
|
||||
@meta_patched_function.register(torch.nn.functional.relu)
|
||||
def torch_nn_func_relu(input, inplace=False):
|
||||
return torch.empty(input.shape, device='meta')
|
||||
return torch.empty(input.shape, device="meta")
|
||||
|
@@ -4,7 +4,7 @@ from ...registry import meta_patched_function
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.matmul)
|
||||
@meta_patched_function.register('matmul') # for built-in op @
|
||||
@meta_patched_function.register("matmul") # for built-in op @
|
||||
def torch_matmul(input, other, *, out=None):
|
||||
# copied from huggingface.utils.fx
|
||||
d1 = input.dim()
|
||||
@@ -44,8 +44,8 @@ def torch_matmul(input, other, *, out=None):
|
||||
|
||||
@meta_patched_function.register(torch.abs)
|
||||
def torch_abs(input, *, out=None):
|
||||
assert out is None, 'out is not supported yet'
|
||||
return torch.empty(input.shape, device='meta')
|
||||
assert out is None, "out is not supported yet"
|
||||
return torch.empty(input.shape, device="meta")
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.bmm)
|
||||
@@ -89,7 +89,7 @@ def torch_addmm(input, mat1, mat2, *, beta=1, alpha=1, out=None):
|
||||
|
||||
@meta_patched_function.register(torch.var_mean)
|
||||
def torch_var_mean(input, dim, unbiased=True, keepdim=False, *, out=None):
|
||||
assert out is None, 'saving to out is not supported yet'
|
||||
var = torch.empty(1).squeeze(0).to('meta')
|
||||
mean = torch.empty(1).squeeze(0).to('meta')
|
||||
assert out is None, "saving to out is not supported yet"
|
||||
var = torch.empty(1).squeeze(0).to("meta")
|
||||
mean = torch.empty(1).squeeze(0).to("meta")
|
||||
return var, mean
|
||||
|
@@ -8,7 +8,6 @@ from ...registry import meta_patched_function
|
||||
|
||||
|
||||
def _ntuple(n, name="parse"):
|
||||
|
||||
def parse(x):
|
||||
if isinstance(x, collections.abc.Iterable):
|
||||
return tuple(x)
|
||||
@@ -24,21 +23,21 @@ _triple = _ntuple(3, "_triple")
|
||||
|
||||
|
||||
def _extract_kwargs(kwargs):
|
||||
if 'stride' in kwargs:
|
||||
stride = kwargs['stride']
|
||||
if "stride" in kwargs:
|
||||
stride = kwargs["stride"]
|
||||
else:
|
||||
stride = 1
|
||||
# TODO: process str type padding
|
||||
if 'padding' in kwargs:
|
||||
padding = kwargs['padding']
|
||||
if "padding" in kwargs:
|
||||
padding = kwargs["padding"]
|
||||
else:
|
||||
padding = 0
|
||||
if 'dilation' in kwargs:
|
||||
dilation = kwargs['dilation']
|
||||
if "dilation" in kwargs:
|
||||
dilation = kwargs["dilation"]
|
||||
else:
|
||||
dilation = 1
|
||||
if 'output_padding' in kwargs:
|
||||
output_padding = kwargs['output_padding']
|
||||
if "output_padding" in kwargs:
|
||||
output_padding = kwargs["output_padding"]
|
||||
else:
|
||||
output_padding = 0
|
||||
|
||||
@@ -61,7 +60,7 @@ def torch_nn_functional_conv1d(input, weight, **kwargs):
|
||||
c_out,
|
||||
l_out,
|
||||
)
|
||||
return torch.empty(result_shape, device='meta')
|
||||
return torch.empty(result_shape, device="meta")
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.nn.functional.conv2d)
|
||||
@@ -82,7 +81,7 @@ def torch_nn_functional_conv2d(input, weight, **kwargs):
|
||||
h_out,
|
||||
w_out,
|
||||
)
|
||||
return torch.empty(result_shape, device='meta')
|
||||
return torch.empty(result_shape, device="meta")
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.nn.functional.conv3d)
|
||||
@@ -105,7 +104,7 @@ def torch_nn_functional_conv3d(input, weight, **kwargs):
|
||||
h_out,
|
||||
w_out,
|
||||
)
|
||||
return torch.empty(result_shape, device='meta')
|
||||
return torch.empty(result_shape, device="meta")
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.nn.functional.conv_transpose1d)
|
||||
@@ -120,13 +119,14 @@ def torch_nn_functional_convtranspose1d(input, weight, **kwargs):
|
||||
kernel_size = weight.shape[2:]
|
||||
l_in = input.shape[-1]
|
||||
c_out = weight.shape[1]
|
||||
l_out = math.floor((l_in - 1) * stride[0] - 2 * padding[0] + dilation[0] * (kernel_size[0] - 1) +
|
||||
output_padding[0] + 1)
|
||||
l_out = math.floor(
|
||||
(l_in - 1) * stride[0] - 2 * padding[0] + dilation[0] * (kernel_size[0] - 1) + output_padding[0] + 1
|
||||
)
|
||||
result_shape = input.shape[:-2] + (
|
||||
c_out,
|
||||
l_out,
|
||||
)
|
||||
return torch.empty(result_shape, device='meta')
|
||||
return torch.empty(result_shape, device="meta")
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.nn.functional.conv_transpose2d)
|
||||
@@ -141,16 +141,18 @@ def torch_nn_functional_convtranspose2d(input, weight, **kwargs):
|
||||
kernel_size = weight.shape[2:]
|
||||
h_in, w_in = input.shape[-2:]
|
||||
c_out = weight.shape[1]
|
||||
h_out = math.floor((h_in - 1) * stride[0] - 2 * padding[0] + dilation[0] * (kernel_size[0] - 1) +
|
||||
output_padding[0] + 1)
|
||||
w_out = math.floor((w_in - 1) * stride[1] - 2 * padding[1] + dilation[1] * (kernel_size[1] - 1) +
|
||||
output_padding[1] + 1)
|
||||
h_out = math.floor(
|
||||
(h_in - 1) * stride[0] - 2 * padding[0] + dilation[0] * (kernel_size[0] - 1) + output_padding[0] + 1
|
||||
)
|
||||
w_out = math.floor(
|
||||
(w_in - 1) * stride[1] - 2 * padding[1] + dilation[1] * (kernel_size[1] - 1) + output_padding[1] + 1
|
||||
)
|
||||
result_shape = input.shape[:-3] + (
|
||||
c_out,
|
||||
h_out,
|
||||
w_out,
|
||||
)
|
||||
return torch.empty(result_shape, device='meta')
|
||||
return torch.empty(result_shape, device="meta")
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.nn.functional.conv_transpose3d)
|
||||
@@ -165,16 +167,19 @@ def torch_nn_functional_convtranspose3d(input, weight, **kwargs):
|
||||
kernel_size = weight.shape[2:]
|
||||
d_in, h_in, w_in = input.shape[-3:]
|
||||
c_out = weight.shape[1]
|
||||
d_out = math.floor((d_in - 1) * stride[0] - 2 * padding[0] + dilation[0] * (kernel_size[0] - 1) +
|
||||
output_padding[0] + 1)
|
||||
h_out = math.floor((h_in - 1) * stride[1] - 2 * padding[1] + dilation[1] * (kernel_size[1] - 1) +
|
||||
output_padding[1] + 1)
|
||||
w_out = math.floor((w_in - 1) * stride[2] - 2 * padding[2] + dilation[2] * (kernel_size[2] - 1) +
|
||||
output_padding[2] + 1)
|
||||
d_out = math.floor(
|
||||
(d_in - 1) * stride[0] - 2 * padding[0] + dilation[0] * (kernel_size[0] - 1) + output_padding[0] + 1
|
||||
)
|
||||
h_out = math.floor(
|
||||
(h_in - 1) * stride[1] - 2 * padding[1] + dilation[1] * (kernel_size[1] - 1) + output_padding[1] + 1
|
||||
)
|
||||
w_out = math.floor(
|
||||
(w_in - 1) * stride[2] - 2 * padding[2] + dilation[2] * (kernel_size[2] - 1) + output_padding[2] + 1
|
||||
)
|
||||
result_shape = input.shape[:-4] + (
|
||||
c_out,
|
||||
d_out,
|
||||
h_out,
|
||||
w_out,
|
||||
)
|
||||
return torch.empty(result_shape, device='meta')
|
||||
return torch.empty(result_shape, device="meta")
|
||||
|
@@ -4,11 +4,7 @@ from ...registry import meta_patched_function
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.nn.functional.embedding)
|
||||
def torch_nn_functional_embedding(input,
|
||||
weight,
|
||||
padding_idx=None,
|
||||
max_norm=None,
|
||||
norm_type=2.0,
|
||||
scale_grad_by_freq=False,
|
||||
sparse=False):
|
||||
def torch_nn_functional_embedding(
|
||||
input, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False
|
||||
):
|
||||
return torch.empty(*input.shape, weight.shape[-1], device="meta")
|
||||
|
@@ -5,16 +5,11 @@ from ...registry import meta_patched_function
|
||||
|
||||
@meta_patched_function.register(torch.nn.functional.layer_norm)
|
||||
def torch_nn_func_layernorm(input, normalized_shape, weight=None, bias=None, eps=1e-05):
|
||||
return torch.empty(input.shape, device='meta')
|
||||
return torch.empty(input.shape, device="meta")
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.nn.functional.batch_norm)
|
||||
def torch_nn_func_batchnorm(input,
|
||||
running_mean,
|
||||
running_var,
|
||||
weight=None,
|
||||
bias=None,
|
||||
training=False,
|
||||
momentum=0.1,
|
||||
eps=1e-05):
|
||||
return torch.empty(input.shape, device='meta')
|
||||
def torch_nn_func_batchnorm(
|
||||
input, running_mean, running_var, weight=None, bias=None, training=False, momentum=0.1, eps=1e-05
|
||||
):
|
||||
return torch.empty(input.shape, device="meta")
|
||||
|
@@ -19,9 +19,9 @@ def operator_getitem(a, b):
|
||||
return t
|
||||
|
||||
def _slice_convert(slice_obj):
|
||||
attrs = {'start': slice_obj.start, 'stop': slice_obj.stop, 'step': slice_obj.step}
|
||||
attrs = {"start": slice_obj.start, "stop": slice_obj.stop, "step": slice_obj.step}
|
||||
new_attrs = _slice_attr_convert(attrs)
|
||||
attr_dict_to_tuple = (new_attrs['start'], new_attrs['stop'], new_attrs['step'])
|
||||
attr_dict_to_tuple = (new_attrs["start"], new_attrs["stop"], new_attrs["step"])
|
||||
return slice(*attr_dict_to_tuple)
|
||||
|
||||
def _slice_attr_convert(attrs):
|
||||
|
@@ -105,14 +105,15 @@ def torch_cat(tensors, dim=None, axis=None, *, out=None):
|
||||
shapes = [t.shape for t in tensors]
|
||||
shape = list(shapes[0])
|
||||
concatenated_dim = sum(shape[dim] for shape in shapes)
|
||||
final_shape = shape[:dim] + [concatenated_dim] + shape[dim + 1:]
|
||||
final_shape = shape[:dim] + [concatenated_dim] + shape[dim + 1 :]
|
||||
return torch.empty(final_shape, device="meta")
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.repeat_interleave)
|
||||
def torch_repeat_interleave(input, repeats, dim=None, output_size=None):
|
||||
assert isinstance(repeats, int) or isinstance(repeats, torch.Tensor), \
|
||||
"Argument 'repeats' should be of type 'torch.Tensor' or 'int'"
|
||||
assert isinstance(repeats, int) or isinstance(
|
||||
repeats, torch.Tensor
|
||||
), "Argument 'repeats' should be of type 'torch.Tensor' or 'int'"
|
||||
|
||||
shape = list(input.shape) if dim is not None else [input.numel()]
|
||||
dim = dim if dim is not None else 0
|
||||
@@ -132,36 +133,36 @@ def torch_tensor_repeat_interleave(self, repeats, dim=None, *, output_size=None)
|
||||
|
||||
@meta_patched_function.register(torch.roll)
|
||||
def torch_roll(input, shifts, dims=None):
|
||||
return torch.empty(input.shape, device='meta')
|
||||
return torch.empty(input.shape, device="meta")
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.full)
|
||||
def torch_full(size, fill_value, *, out=None, dtype=None, layout=torch.strided, device=None, requires_grad=False):
|
||||
assert out is None, 'assigning result to out is not supported yet'
|
||||
return torch.empty(size, device='meta', dtype=dtype, layout=layout, requires_grad=requires_grad)
|
||||
assert out is None, "assigning result to out is not supported yet"
|
||||
return torch.empty(size, device="meta", dtype=dtype, layout=layout, requires_grad=requires_grad)
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.max)
|
||||
def torch_max(input, dim=None, keepdim=False, *, out=None):
|
||||
assert out is None, 'assigning value to out is not supported yet'
|
||||
assert out is None, "assigning value to out is not supported yet"
|
||||
if dim is not None:
|
||||
if isinstance(dim, int):
|
||||
shape = list(input.shape)
|
||||
shape.pop(dim)
|
||||
if keepdim:
|
||||
shape.insert(dim, 1)
|
||||
return torch.empty(shape, device='meta', dtype=input.dtype), torch.empty(shape,
|
||||
device='meta',
|
||||
dtype=input.dtype)
|
||||
return torch.empty(shape, device="meta", dtype=input.dtype), torch.empty(
|
||||
shape, device="meta", dtype=input.dtype
|
||||
)
|
||||
elif isinstance(dim, torch.Tensor):
|
||||
# when dim is a 0D or 1D tensor, it will maintain the same shape
|
||||
num_dims = dim.dim()
|
||||
if num_dims in [0, 1]:
|
||||
return torch.empty_like(input, device='meta')
|
||||
return torch.empty_like(input, device="meta")
|
||||
else:
|
||||
raise ValueError(f"Expected dim to a 0D or 1D tensor but got {num_dims} dimensions")
|
||||
else:
|
||||
return torch.empty([], device='meta', dtype=input.dtype)
|
||||
return torch.empty([], device="meta", dtype=input.dtype)
|
||||
|
||||
|
||||
@meta_patched_function.register(torch.Tensor.cpu)
|
||||
|
Reference in New Issue
Block a user