形状和布局

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 isminor”。

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

a b c d e f

对于 N 维度数组,从次要到主要维度顺序(从 N-10)类似于“行主”(排名为 2)。假设维度的单调排序,我们在代码中引用此布局的另一种方式是“dim 0 is main”。

默认的从小到大排序

新创建的形状的默认布局为“维度顺序为从主到次要”(类似于排名为 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