jax.lax.collapse

jax.lax.collapse(operand, start_dimension, stop_dimension=None)

将数组的维度折叠成单个维度。

例如，如果 operand 是一个形状为 [2, 3, 4] 的数组，则 collapse(operand, 0, 2).shape == [6, 4]。折叠维度的元素以主到次的方式排列，即编号最小的维度变化最慢。

参数:

• operand (数组) – 输入数组。

• start_dimension (int) – 要折叠的维度的起始位置（包含）。

• stop_dimension (int | None | None) – 要折叠的维度的结束位置（不包含）。传递 None 以折叠起始位置之后的所有维度。

返回:

一个数组，其中维度 [start_dimension, stop_dimension) 已被折叠（扁平化）成单个维度。

返回类型:

数组