XLA: بهینه سازی کامپایلر برای یادگیری ماشین

API کامپایل صریح یک کنترل دقیق برای انتخاب توابع کامپایل ارائه می دهد. به عنوان مثال، تابع TensorFlow زیر که آموزش MNIST را انجام می دهد با XLA کامپایل شده است:

@tf.function(jit_compile=True)
def train_mnist(images, labels):
    images, labels = cast(images, labels)

    with tf.GradientTape() as tape:
      predicted_labels = layer(images)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=predicted_labels, labels=labels
      ))
    layer_variables = layer.trainable_variables
    grads = tape.gradient(loss, layer_variables)
    optimizer.apply_gradients(zip(grads, layer_variables))

API jit_compile دارای معنایی باید کامپایل باشد : یا کل تابع با XLA کامپایل شده است یا یک errors.InvalidArgumentError استثنا پرتاب می شود. XLA در حال حاضر نمی تواند توابعی را در جایی که ابعاد قابل استنباط نیست کامپایل کند: یعنی اگر امکان استنتاج ابعاد همه تانسورها بدون اجرای کل محاسبات وجود نداشته باشد. به عنوان مثال، تابع زیر کامپایل نمی شود:

@tf.function
def not_compilable(x):
  return tf.unique(x)

شکل ها می توانند در طول اجراها متفاوت باشند:

@tf.function(jit_compile=True)
def recompiled_on_launch(a, b):
  return a + b

recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))

برای مثال استفاده دقیق تر، و یک ویدیوی آموزشی در مورد jit_compile=True به colab آموزشی مراجعه کنید.

model.compile(optimizer="adam", jit_compile=True)

@tf.function(jit_compile=True)
def step_fn():
  t = tf.ones(shape=[100], dtype=tf.float32)
  ctx = tf.distribute.get_replica_context()
  return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)

@tf.function
def run_fn():
  return strategy.run(step_fn)

$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program

$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

$ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
  TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
  XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
    my/tensorflow/program"