jax.lax.collapse#

jax.lax.collapse(operand, start_dimension, stop_dimension=None)[来源]#

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

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

参数:
  • operand (Array) – 一个输入数组。

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

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

返回:

一个数组,其中维度 [start_dimension, stop_dimension) 已折叠(扁平化)为单个维度。

返回类型:

Array