提交 c70a8869 authored 作者: Ricardo Vieira's avatar Ricardo Vieira 提交者: Ricardo Vieira

Stack eager optimization for single tensor

上级 0de0fa9b
......@@ -2943,6 +2943,8 @@ def stack(tensors: Sequence["TensorLike"], axis: int = 0):
):
# In case there is direct scalar
tensors = list(map(as_tensor_variable, tensors))
if len(tensors) == 1:
return atleast_1d(tensors[0])
dtype = ps.upcast(*[i.dtype for i in tensors])
return MakeVector(dtype)(*tensors)
return join(axis, *[shape_padaxis(t, axis) for t in tensors])
......
Markdown 格式
0%
您添加了 0 到此讨论。请谨慎行事。
请先完成此评论的编辑!
注册 或者 后发表评论