シェイプとレイアウト

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 次元の配列には多くの場合、次元に関連付けられた特定の文字があります。2 次元配列の例:

    • ディメンション 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 プロトコルは、配列がメモリ内でどのように表現されるかを記述します。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 次元配列の場合は 0 から N-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）に似ています。ディメンションの順序が単調であると仮定すると、コードでこのレイアウトを単に「暗 0 はマイナー」と呼ぶこともできます。

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

a b c d e f

ディメンション配列のN-1 から 0 のマイナーからメジャーへの順は、N ディメンション配列の行メジャー（ランク 2）に似ています。ディメンションの順序が単調であると仮定すると、コード内でこのレイアウトを参照する別の方法としては、「dim 0 is large（ディメンション 0 はメジャー）」というものもあります。

デフォルトのマイナーから優先への順序

新しく作成されたシェイプのデフォルトのレイアウトは、「ディメンションの順序はメジャーからマイナー」です（ランク 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.h と layout_util.h をご覧ください。