形状とレイアウト

XLA shape proto（xla_data.proto）は、N 次元配列（略して配列）の階数とサイズ、データ型について記述します。

用語、表記、規則

  • 配列の階数は次元の数と同じです。配列の実際の階数は、1 より大きいサイズを持つ次元数です。

  • 次元は、N 次元配列の場合、0 から N-1 で番号付けされます。次元番号は、便宜的につけられたラベルです。次元の順序は、形状のレイアウトの優先順序（minor/major order）を意味しているわけではありません。レイアウトは、Layout proto によって決まります。

  • 規則として、次元は次元番号の昇順にリスト化されます。たとえば、サイズが [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

    3 次元配列の場合は以下のようになります。

    • 次元 0: z
    • 次元1: y
    • 次元 2: x

    4 次元配列では以下のようになります。

    • 次元 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 次元の配列では 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 の場合）。モノトニックな次元順序を仮定した場合、コード内でこのレイアウトを参照するために使用できる別名は、単純に「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_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 をご覧ください。