シェイプとレイアウト

XLA Shape proto(xla_data.proto)は、N 次元配列(簡潔では配列)のランク、サイズ、データ型を記述します。

用語、表記、規則

  • 配列の階数は次元の数に等しくなります。配列の真のランクは、サイズが 1 より大きい次元の数です。

  • N ディメンション配列の場合、ディメンションには 0 から N-1 まで番号が付けられます。ディメンション番号は、便宜上任意のラベルです。寸法番号の順序は、シェイプのレイアウトにおける特定のマイナーまたはメジャーな順序を示すものではありません。レイアウトは Layout プロトコルによって決まります。

  • 慣例として、ディメンションはディメンション番号の昇順で記載します。たとえば、サイズが [A x B x C] の 3 次元配列の場合、ディメンション 0 のサイズは A、ディメンション 1 のサイズは B、ディメンション 2 のサイズは C です。

    XLA の一部のユーティリティは Python に似た負のインデックスもサポートしています。ディメンション -1 は最後のディメンションです(N ディメンション配列の N-1 に相当します)。たとえば、上記の 3 次元配列の場合、ディメンション -1 のサイズは C、ディメンション -2 のサイズは B となります。

  • 2 次元、3 次元、4 次元の配列では多くの場合、次元に関連付けられた特定の文字があります。たとえば、2D 配列の場合は次のようになります。

    • ディメンション 0: y
    • ディメンション 1: x

    3D 配列の場合:

    • ディメンション 0: z
    • ディメンション 1: y
    • ディメンション 2: x

    4D 配列の場合:

    • ディメンション 0: p
    • ディメンション 1: z
    • ディメンション 2: y
    • ディメンション 3: x
  • ディメンションを受け取る XLA API の関数は、ディメンション番号の昇順で受け取ります。これは、ディメンションを initializer_list として渡すときに使用される順序と一致します。たとえば、

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

    ディメンション サイズの配列が [A, B, C, D] のシーケンスで構成されるシェイプが作成されます。

レイアウト

Layout プロトコルは、配列がメモリでどのように表されるかを記述します。Layout プロトコルには次のフィールドが含まれています。

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] の 2 次元配列があるとします。

a b c
d e f

ここで、ディメンション 0 はサイズ 2、ディメンション 1 はサイズ 3 です。レイアウト内の minor_to_major フィールドが [0, 1] の場合、ディメンション 0 が最小のディメンションで、ディメンション 1 が最上位のディメンションです。これは、リニアメモリ内の次のレイアウトに対応します。

a d b e c f

0 から N-1 までのこのマイナーからメジャーへのディメンションの順序は、列メジャー(ランク 2)に似ています。寸法が単調であると仮定すると、コード内でこのレイアウトを指す別の言い方をすると、単に「dim 0 はマイナー」です。

一方、レイアウト内の minor_to_major フィールドが [1, 0] の場合、リニアメモリ内のレイアウトは次のようになります。

a b c d e f

N ディメンション配列の N-1 から 0 へのマイナーからメジャーへのディメンションの順序は、行優先(ランク 2)に似ています。次元が単調であると仮定すると、コードでこのレイアウトを参照する別の方法は、単に「dim 0 isMajor」です。

デフォルトのマイナーからメジャーへの順序

新しく作成された Shapes のデフォルト レイアウトは、「ディメンションの順序はメジャーからマイナーです」(ランク 2 の行優先と同じ)です。

パディング

パディングは、オプションの padded_dimensions フィールドと padding_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 をご覧ください。