Save the date! Google I/O returns May 18-20 Register now

形状和布局

XLA Shape proto (xla_data.proto) 描述了 N 维数组(简称数组)的秩、大小和数据类型。

术语、表示法和惯例

  • 数组的秩等于维数。数组的真实秩等于大小大于 1 的维度的数量。

  • 对于 N 维数组,维度被编号为 0N-1。为方便起见,维度编号为任意标签。这些维度编号的顺序并不表示形状布局中的特定主次顺序。布局由 Layout proto 确定。

  • 按照惯例,维度将按照维度编号的升序排列。例如,对于大小为 [A x B x C] 的三维数组,维度 0 的大小为 A、维度 1 的大小为 B、维度 2 的大小为 C

    XLA 中的某些效用函数还支持负索引,类似于 Python;维度 -1 为最后一个维度(等同于 N 维数组中的维度 N-1)。以上文中的三维数组为例,维度 -1 的大小为 C、维度 -2 的大小为 B,依此类推。

  • 二维、三维和四维数组通常具有与维度关联的特定字母。例如,对于二维数组:

    • 维度 0:y
    • 维度 1:x

    对于三维数组:

    • 维度 0:z
    • 维度 1:y
    • 维度 2:x

    对于四维数组:

    • 维度 0:p
    • 维度 1:z
    • 维度 2:y
    • 维度 3:x
  • XLA API 中接受维度的函数可利用这些字母实现维度编号的升序排列。这与将维度作为 initializer_list 传递时使用的顺序相符;例如:

    ShapeUtil::MakeShape(F32, {A, B, C, D})

    将创建一个形状,其维度大小数组由序列 [A, B, C, D] 组成。

布局

Layout proto 描述了数组在内存中的表示方式。Layout proto 包含以下字段:

message Layout {
  repeated int64 minor_to_major = 1;
  repeated int64 padded_dimensions = 2;
  optional PaddingValue padding_value = 3;
}

从小到大的维度排序

minor_to_major 是唯一的必需字段。此字段描述了形状中从小到大的维度排序。minor_to_major 中的值为数组的维度排序(N 维数组的排序为 0N-1),第一个值为最小维度,最后一个值为最大维度。最小维度是在遍历线性内存中排列的数组元素时变化最快的维度。

例如,请思考以下大小为 [2 x 3] 的二维数组:

a b c
d e f

在该数组中,维度 0 的大小为 2、维度 1 的大小为 3。如果布局中的 minor_to_major 字段为 [0, 1],那么维度 0 为最小维度,维度 1 为最大维度。这对应于线性内存中的以下布局:

a d b e c f

0N-1 的从小到大的维度排序类似于列优先(秩为 2 时)。假定维度符合单调排序,那么我们在代码中引用此布局时可以使用另一种简单名称“dim 0 is minor”。

另一方面,如果布局中的 minor_to_major 字段为 [1, 0],则线性内存中的布局为:

a b c d e f

N 维数组中,N-1 降序至 0 的从大到小的维度排序类似于行优先(秩为 2 时)。假定维度符合单调排序,那么我们在代码中引用此布局时可以使用另一种简单名称“dim 0 is major”。

默认的从大到小的排序

新创建形状的默认布局为“从大到小的维度排序”(类似于秩为 2 时的行优先)。

填充

填充在可选的 padded_dimensionspadding_value 字段中进行定义。padded_dimensions 字段描述了每个维度填充至的大小(宽度)。如果存在,padded_dimensions 中的元素数量必须与形状的秩相等。

以上文中的 [2 x 3] 数组为例,如果 padded_dimensions[3, 5],那么维度 0 的宽度将填充至 3,维度 1 的宽度填充至 5。线性内存中的布局(假定填充值为 0,并采用列优先布局)为:

a d 0 b e 0 c f 0 0 0 0 0 0 0

这等效于采用相同从小到大的维度排序的以下数组的布局:

a b c 0 0
d e f 0 0
0 0 0 0 0

索引到数组

index_util.h 中的 IndexUtil 类提供了用于在给定形状和布局的情况下在多维索引和线性索引之间转换的效用函数。多维索引包含每个维度的 int64 索引。线性索引为单个 int64 值,可索引到存储数组的缓冲区。请参阅同一目录中的 shape_util.hlayout_util.h,获取能够简化形状和布局创建与操作的效用函数。