import paddle
n = 44 * 1024  # 需要34G显存
x = paddle.randn([n, n], dtype='float16')
index = paddle.stack([
    paddle.arange(n, dtype='int32').reshape([n, 1]).repeat_interleave(n, 1),
    paddle.arange(n, dtype='int32').reshape([1, n]).repeat_interleave(n, 0),
], axis=2)
y = paddle.gather_nd(x, index)
del index
assert paddle.sum(x == y) == n * n  # 如果gather_nd实现正确，x和y应该是相等的

import paddle
n = 44 * 1024  # 需要34G显存
x = paddle.randn([n, n], dtype='float16')
index = paddle.stack([
    paddle.arange(n, dtype='int32').reshape([n, 1]).repeat_interleave(n, 1),
    paddle.arange(n, dtype='int32').reshape([1, n]).repeat_interleave(n, 0),
], axis=2)
y = paddle.scatter_nd(index, x, [n, n])
del index
assert paddle.sum(x == y) == n * n  # 如果scatter_nd实现正确，x和y应该是相等的