未验证 提交 34b6860e 编写于 作者: F Feiyu Chan 提交者: GitHub

fix fftshift/ifftshift on static mode (#36748)

* fix fftshift/ifftshift on static mode
* update roll_op version
* add more test cases for fftshift/ifftshift
上级 6838a187
......@@ -183,7 +183,12 @@ REGISTER_OP_VERSION(roll)
"(std::vector<int64_t>) Axis along which to roll. "
"It must have the same size with shifts, or size = 0.",
std::vector<int64_t>())
.DeleteAttr(
"dims",
.DeleteAttr("dims",
"(std::vector<int64_t>) Dims along which to roll. "
"It must have the same size with shifts, or size = 0."));
"It must have the same size with shifts, or size = 0."))
.AddCheckpoint(
R"ROC(Upgrade roll add a dispensable input "ShiftsTensor".)ROC",
paddle::framework::compatible::OpVersionDesc().NewInput(
"ShiftsTensor",
"The number of places by which the elements of"
"the tensor are shifted."));
......@@ -1300,13 +1300,13 @@ def fftshift(x, axes=None, name=None):
shape = paddle.shape(x)
if axes is None:
# shift all axes
rank = paddle.rank(x).reshape([1])
axes = axes or paddle.arange(0, rank)
shifts = [size // 2 for size in shape]
rank = len(x.shape)
axes = list(range(0, rank))
shifts = shape // 2
elif isinstance(axes, int):
shifts = shape[axes] // 2
else:
shifts = [shape[ax] // 2 for ax in axes]
shifts = paddle.concat([shape[ax] // 2 for ax in axes])
return paddle.roll(x, shifts, axes, name=name)
......@@ -1343,13 +1343,13 @@ def ifftshift(x, axes=None, name=None):
shape = paddle.shape(x)
if axes is None:
# shift all axes
rank = paddle.rank(x).reshape([1])
axes = axes or paddle.arange(0, rank)
shifts = [-size // 2 for size in shape]
rank = len(x.shape)
axes = list(range(0, rank))
shifts = shape // 2
elif isinstance(axes, int):
shifts = -shape[axes] // 2
else:
shifts = [-shape[ax] // 2 for ax in axes]
shifts = paddle.concat([-shape[ax] // 2 for ax in axes])
return paddle.roll(x, shifts, axes, name=name)
......
......@@ -1009,10 +1009,11 @@ class TestRfftFreq(unittest.TestCase):
@place(DEVICES)
@parameterize((TEST_CASE_NAME, 'x', 'axes', 'dtype'), [
('test_1d', np.random.randn(10), (0, ), 'float64'),
@parameterize(
(TEST_CASE_NAME, 'x', 'axes', 'dtype'),
[('test_1d', np.random.randn(10), (0, ), 'float64'),
('test_2d', np.random.randn(10, 10), (0, 1), 'float64'),
])
('test_2d_with_all_axes', np.random.randn(10, 10), None, 'float64')])
class TestFftShift(unittest.TestCase):
def test_fftshift(self):
"""Test fftshift with norm condition
......@@ -1030,6 +1031,7 @@ class TestFftShift(unittest.TestCase):
@parameterize((TEST_CASE_NAME, 'x', 'axes'), [
('test_1d', np.random.randn(10), (0, ), 'float64'),
('test_2d', np.random.randn(10, 10), (0, 1), 'float64'),
('test_2d_with_all_axes', np.random.randn(10, 10), None, 'float64'),
])
class TestIfftShift(unittest.TestCase):
def test_ifftshift(self):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册