jax.lax.reshape

目录

jax.lax.reshape#

jax.lax.reshape(operand, new_sizes, dimensions=None)[源代码][源代码]#

封装了XLA的 Reshape 操作符。

对于插入/移除大小为1的维度，建议使用 lax.squeeze / lax.expand_dims。这些操作保留了可能对高级变换规则有用的轴标识信息。

参数:

  • operand (ArrayLike) – 要重塑的数组。

  • new_sizes (Shape) – 指定结果形状的整数序列。最终数组的大小必须与输入的大小匹配。

  • dimensions (Sequence[int] | None) – 可选的整数序列，指定输入形状的排列顺序。如果指定，长度必须与 operand.shape 匹配。

返回:

重塑后的数组。

返回类型:

out

示例

从一维到二维的简单重塑：

>>> x = jnp.arange(6)
>>> y = reshape(x, (2, 3))
>>> y
Array([[0, 1, 2],
             [3, 4, 5]], dtype=int32)

重塑回一维：

>>> reshape(y, (6,))
Array([0, 1, 2, 3, 4, 5], dtype=int32)

通过维度置换重塑为一维：

>>> reshape(y, (6,), (1, 0))
Array([0, 3, 1, 4, 2, 5], dtype=int32)
目录