Используйте XLA с функцией tf.

Посмотреть на TensorFlow.org

Запускаем в Google Colab

Скачать блокнот

Посмотреть исходный код на GitHub

Этот учебник обучает модель TensorFlow классифицировать набор данных MNIST, где обучающая функция компилируется с использованием XLA.

Сначала загрузите TensorFlow и включите активное выполнение.

import tensorflow as tf
tf.compat.v1.enable_eager_execution()

Затем определите некоторые необходимые константы и подготовьте набор данных MNIST.

# Size of each input image, 28 x 28 pixels
IMAGE_SIZE = 28 * 28
# Number of distinct number labels, [0..9]
NUM_CLASSES = 10
# Number of examples in each training batch (step)
TRAIN_BATCH_SIZE = 100
# Number of training steps to run
TRAIN_STEPS = 1000

# Loads MNIST dataset.
train, test = tf.keras.datasets.mnist.load_data()
train_ds = tf.data.Dataset.from_tensor_slices(train).batch(TRAIN_BATCH_SIZE).repeat()

# Casting from raw data to the required datatypes.
def cast(images, labels):
  images = tf.cast(
      tf.reshape(images, [-1, IMAGE_SIZE]), tf.float32)
  labels = tf.cast(labels, tf.int64)
  return (images, labels)

Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11493376/11490434 [==============================] - 0s 0us/step
11501568/11490434 [==============================] - 0s 0us/step

Наконец, определите модель и оптимизатор. В модели используется один плотный слой.

layer = tf.keras.layers.Dense(NUM_CLASSES)
optimizer = tf.keras.optimizers.Adam()

Определите функцию обучения

В функции обучения вы получаете предсказанные метки, используя слой, определенный выше, а затем минимизируете градиент потерь с помощью оптимизатора. Для того , чтобы скомпилировать вычисления с использованием XLA, поместите его в tf.function с jit_compile=True .

@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))

Обучите и протестируйте модель

После того, как вы определили функцию обучения, определите модель.

for images, labels in train_ds:
  if optimizer.iterations > TRAIN_STEPS:
    break
  train_mnist(images, labels)

И, наконец, проверяем точность:

images, labels = cast(test[0], test[1])
predicted_labels = layer(images)
correct_prediction = tf.equal(tf.argmax(predicted_labels, 1), labels)
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))
print("Prediction accuracy after training: %s" % accuracy)

Prediction accuracy after training: tf.Tensor(0.8721, shape=(), dtype=float32)

За кулисами компилятор XLA скомпилировал всю функцию TF в HLO, что позволило провести оптимизацию слияния. Используя интроспекцию объектов, мы можем увидеть код Hlo (других интересных возможных значения для «стадии» являются optimized_hlo для ПЗА после оптимизаций и optimized_hlo_dot для графа Graphviz):

print(train_mnist.experimental_get_compiler_ir(images, labels)(stage='hlo'))

HloModule a_inference_train_mnist_5284__.204, input_output_alias={ {0}: (2, {}, may-alias), {1}: (3, {}, may-alias), {2}: (5, {}, may-alias), {3}: (8, {}, may-alias), {4}: (9, {}, may-alias), {5}: (10, {}, may-alias), {6}: (11, {}, may-alias) }

%max_float_.65 (x.66: f32[], y.67: f32[]) -> f32[] {
  %x.66 = f32[] parameter(0)
  %y.67 = f32[] parameter(1)
  ROOT %maximum.68 = f32[] maximum(f32[] %x.66, f32[] %y.67)
}

%add_float_.75 (x.76: f32[], y.77: f32[]) -> f32[] {
  %x.76 = f32[] parameter(0)
  %y.77 = f32[] parameter(1)
  ROOT %add.78 = f32[] add(f32[] %x.76, f32[] %y.77)
}

%add_float_.94 (x.95: f32[], y.96: f32[]) -> f32[] {
  %x.95 = f32[] parameter(0)
  %y.96 = f32[] parameter(1)
  ROOT %add.97 = f32[] add(f32[] %x.95, f32[] %y.96)
}

%Mean-reduction.106 (x.107: f32[], y.108: f32[]) -> f32[] {
  %x.107 = f32[] parameter(0)
  %y.108 = f32[] parameter(1)
  ROOT %add.109 = f32[] add(f32[] %x.107, f32[] %y.108)
}

%region_0.123 (Arg_0.124: f32[], Arg_1.125: f32[]) -> f32[] {
  %Arg_0.124 = f32[] parameter(0)
  %Arg_1.125 = f32[] parameter(1)
  ROOT %add.126 = f32[] add(f32[] %Arg_0.124, f32[] %Arg_1.125), metadata={op_name="gradient_tape/dense/BiasAdd/BiasAddGrad"}
}

ENTRY %a_inference_train_mnist_5284__.204 (arg0.1: f32[10000,784], arg1.2: s64[10000], arg2.3: f32[784,10], arg3.4: f32[10], arg4.5: f32[], arg5.6: s64[], arg6.7: f32[], arg7.8: f32[], arg8.9: f32[784,10], arg9.10: f32[784,10], arg10.11: f32[10], arg11.12: f32[10]) -> (f32[784,10], f32[10], s64[], f32[784,10], f32[784,10], /*index=5*/f32[10], f32[10]) {
  %constant.15 = f32[] constant(1), metadata={op_type="Sub" op_name="Adam/sub_2" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=141}
  %constant.16 = f32[] constant(1), metadata={op_type="Sub" op_name="Adam/sub_2" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=141}
  %arg6.7 = f32[] parameter(6), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %subtract.17 = f32[] subtract(f32[] %constant.16, f32[] %arg6.7), metadata={op_type="Sub" op_name="Adam/sub_2" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=141}
  %constant.18 = f32[] constant(1), metadata={op_type="Sub" op_name="Adam/sub_3" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=144}
  %constant.19 = f32[] constant(1), metadata={op_type="Sub" op_name="Adam/sub_3" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=144}
  %arg7.8 = f32[] parameter(7), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %subtract.20 = f32[] subtract(f32[] %constant.19, f32[] %arg7.8), metadata={op_type="Sub" op_name="Adam/sub_3" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=144}
  %constant.21 = s64[] constant(1), metadata={op_type="AddV2" op_name="Adam/add" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=127}
  %constant.26 = f32[] constant(1), metadata={op_type="Sub" op_name="Adam/sub_1" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=133}
  %constant.30 = f32[] constant(1), metadata={op_type="Sub" op_name="Adam/sub" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=133}
  %arg4.5 = f32[] parameter(4), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %constant.31 = f32[] constant(1), metadata={op_type="Sub" op_name="Adam/sub" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=133}
  %arg5.6 = s64[] parameter(5), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %constant.22 = s64[] constant(1), metadata={op_type="AddV2" op_name="Adam/add" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=127}
  %add.23 = s64[] add(s64[] %arg5.6, s64[] %constant.22), metadata={op_type="AddV2" op_name="Adam/add" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=127}
  %convert.24 = f32[] convert(s64[] %add.23), metadata={op_type="Cast" op_name="Adam/Cast_1" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=127}
  %power.29 = f32[] power(f32[] %arg7.8, f32[] %convert.24), metadata={op_type="Pow" op_name="Adam/Pow_1" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=131}
  %subtract.32 = f32[] subtract(f32[] %constant.31, f32[] %power.29), metadata={op_type="Sub" op_name="Adam/sub" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=133}
  %sqrt.33 = f32[] sqrt(f32[] %subtract.32), metadata={op_name="Adam/Sqrt"}
  %constant.27 = f32[] constant(1), metadata={op_type="Sub" op_name="Adam/sub_1" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=133}
  %power.25 = f32[] power(f32[] %arg6.7, f32[] %convert.24), metadata={op_type="Pow" op_name="Adam/Pow" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=130}
  %subtract.28 = f32[] subtract(f32[] %constant.27, f32[] %power.25), metadata={op_type="Sub" op_name="Adam/sub_1" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=133}
  %divide.34 = f32[] divide(f32[] %sqrt.33, f32[] %subtract.28), metadata={op_type="RealDiv" op_name="Adam/truediv" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=133}
  %multiply.35 = f32[] multiply(f32[] %arg4.5, f32[] %divide.34), metadata={op_type="Mul" op_name="Adam/mul" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=133}
  %constant.36 = s32[2]{0} constant({-1, 784}), metadata={op_type="Reshape" op_name="Reshape" source_file="/tmp/ipykernel_31814/494562224.py" source_line=16}
  %arg1.2 = s64[10000]{0} parameter(1), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %reshape.14 = s64[10000]{0} reshape(s64[10000]{0} %arg1.2)
  %broadcast.47 = s64[10000,10]{1,0} broadcast(s64[10000]{0} %reshape.14), dimensions={0}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %iota.46 = s64[10000,10]{1,0} iota(), iota_dimension=1, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %compare.48 = pred[10000,10]{1,0} compare(s64[10000,10]{1,0} %broadcast.47, s64[10000,10]{1,0} %iota.46), direction=EQ, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.43 = f32[] constant(1), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.45 = f32[10000,10]{1,0} broadcast(f32[] %constant.43), dimensions={}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.42 = f32[] constant(0), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.44 = f32[10000,10]{1,0} broadcast(f32[] %constant.42), dimensions={}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %select.49 = f32[10000,10]{1,0} select(pred[10000,10]{1,0} %compare.48, f32[10000,10]{1,0} %broadcast.45, f32[10000,10]{1,0} %broadcast.44), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.57 = s64[] constant(0), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.58 = s64[10000]{0} broadcast(s64[] %constant.57), dimensions={}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %compare.59 = pred[10000]{0} compare(s64[10000]{0} %broadcast.58, s64[10000]{0} %reshape.14), direction=LE, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.54 = s64[] constant(10), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.55 = s64[10000]{0} broadcast(s64[] %constant.54), dimensions={}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %compare.56 = pred[10000]{0} compare(s64[10000]{0} %reshape.14, s64[10000]{0} %broadcast.55), direction=LT, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %and.60 = pred[10000]{0} and(pred[10000]{0} %compare.59, pred[10000]{0} %compare.56), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.52 = f32[] constant(0), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.53 = f32[10000]{0} broadcast(f32[] %constant.52), dimensions={}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.50 = f32[] constant(nan), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.51 = f32[10000]{0} broadcast(f32[] %constant.50), dimensions={}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %select.61 = f32[10000]{0} select(pred[10000]{0} %and.60, f32[10000]{0} %broadcast.53, f32[10000]{0} %broadcast.51), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.62 = f32[10000,10]{1,0} broadcast(f32[10000]{0} %select.61), dimensions={0}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %add.63 = f32[10000,10]{1,0} add(f32[10000,10]{1,0} %select.49, f32[10000,10]{1,0} %broadcast.62), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %negate.90 = f32[10000,10]{1,0} negate(f32[10000,10]{1,0} %add.63), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.86 = f32[] constant(0), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.87 = f32[10000,10]{1,0} broadcast(f32[] %constant.86), dimensions={}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %compare.88 = pred[10000,10]{1,0} compare(f32[10000,10]{1,0} %add.63, f32[10000,10]{1,0} %broadcast.87), direction=EQ, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.84 = f32[] constant(0), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.85 = f32[10000,10]{1,0} broadcast(f32[] %constant.84), dimensions={}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %arg0.1 = f32[10000,784]{1,0} parameter(0), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %reshape.13 = f32[10000,784]{1,0} reshape(f32[10000,784]{1,0} %arg0.1)
  %reshape.37 = f32[10000,784]{1,0} reshape(f32[10000,784]{1,0} %reshape.13), metadata={op_type="Reshape" op_name="Reshape" source_file="/tmp/ipykernel_31814/494562224.py" source_line=16}
  %arg2.3 = f32[784,10]{1,0} parameter(2), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %dot.38 = f32[10000,10]{1,0} dot(f32[10000,784]{1,0} %reshape.37, f32[784,10]{1,0} %arg2.3), lhs_contracting_dims={1}, rhs_contracting_dims={0}, metadata={op_type="MatMul" op_name="dense/MatMul" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/layers/core.py" source_line=1229}
  %transpose.39 = f32[10000,10]{1,0} transpose(f32[10000,10]{1,0} %dot.38), dimensions={0,1}, metadata={op_type="MatMul" op_name="dense/MatMul" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/layers/core.py" source_line=1229}
  %arg3.4 = f32[10]{0} parameter(3), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %broadcast.40 = f32[10000,10]{1,0} broadcast(f32[10]{0} %arg3.4), dimensions={1}, metadata={op_type="BiasAdd" op_name="dense/BiasAdd" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/layers/core.py" source_line=1240}
  %add.41 = f32[10000,10]{1,0} add(f32[10000,10]{1,0} %transpose.39, f32[10000,10]{1,0} %broadcast.40), metadata={op_type="BiasAdd" op_name="dense/BiasAdd" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/layers/core.py" source_line=1240}
  %constant.64 = f32[] constant(-inf), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %reduce.69 = f32[10000]{0} reduce(f32[10000,10]{1,0} %add.41, f32[] %constant.64), dimensions={1}, to_apply=%max_float_.65, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.70 = f32[10000,10]{1,0} broadcast(f32[10000]{0} %reduce.69), dimensions={0}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %subtract.71 = f32[10000,10]{1,0} subtract(f32[10000,10]{1,0} %add.41, f32[10000,10]{1,0} %broadcast.70), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %exponential.72 = f32[10000,10]{1,0} exponential(f32[10000,10]{1,0} %subtract.71), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %convert.73 = f32[10000,10]{1,0} convert(f32[10000,10]{1,0} %exponential.72), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.74 = f32[] constant(0), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %reduce.79 = f32[10000]{0} reduce(f32[10000,10]{1,0} %convert.73, f32[] %constant.74), dimensions={1}, to_apply=%add_float_.75, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %convert.80 = f32[10000]{0} convert(f32[10000]{0} %reduce.79), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %log.81 = f32[10000]{0} log(f32[10000]{0} %convert.80), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %broadcast.82 = f32[10000,10]{1,0} broadcast(f32[10000]{0} %log.81), dimensions={0}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %subtract.83 = f32[10000,10]{1,0} subtract(f32[10000,10]{1,0} %subtract.71, f32[10000,10]{1,0} %broadcast.82), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %select.89 = f32[10000,10]{1,0} select(pred[10000,10]{1,0} %compare.88, f32[10000,10]{1,0} %broadcast.85, f32[10000,10]{1,0} %subtract.83), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %multiply.91 = f32[10000,10]{1,0} multiply(f32[10000,10]{1,0} %negate.90, f32[10000,10]{1,0} %select.89), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %convert.93 = f32[10000,10]{1,0} convert(f32[10000,10]{1,0} %multiply.91), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.92 = f32[] constant(0), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %reduce.98 = f32[10000]{0} reduce(f32[10000,10]{1,0} %convert.93, f32[] %constant.92), dimensions={1}, to_apply=%add_float_.94, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %convert.99 = f32[10000]{0} convert(f32[10000]{0} %reduce.98), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %convert.103 = f32[10000]{0} convert(f32[10000]{0} %convert.99), metadata={op_type="Mean" op_name="Mean" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.104 = f32[] constant(0), metadata={op_type="Mean" op_name="Mean" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %convert.105 = f32[] convert(f32[] %constant.104), metadata={op_type="Mean" op_name="Mean" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %reduce.110 = f32[] reduce(f32[10000]{0} %convert.103, f32[] %convert.105), dimensions={0}, to_apply=%Mean-reduction.106, metadata={op_type="Mean" op_name="Mean" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.111 = s32[] constant(10000), metadata={op_type="Mean" op_name="Mean" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %convert.112 = f32[] convert(s32[] %constant.111), metadata={op_type="Mean" op_name="Mean" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %divide.113 = f32[] divide(f32[] %reduce.110, f32[] %convert.112), metadata={op_type="Mean" op_name="Mean" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %convert.114 = f32[] convert(f32[] %divide.113), metadata={op_type="Mean" op_name="Mean" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %constant.115 = f32[] constant(0.0001), metadata={op_type="Mul" op_name="gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=11}
  %broadcast.116 = f32[10000,1]{1,0} broadcast(f32[] %constant.115), dimensions={}, metadata={op_type="Mul" op_name="gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=11}
  %arg8.9 = f32[784,10]{1,0} parameter(8), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %constant.117 = f32[] constant(0.0001), metadata={op_type="Mul" op_name="gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=11}
  %broadcast.118 = f32[10000,1]{1,0} broadcast(f32[] %constant.117), dimensions={}, metadata={op_type="Mul" op_name="gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=11}
  %reshape.119 = f32[10000]{0} reshape(f32[10000,1]{1,0} %broadcast.118), metadata={op_type="Mul" op_name="gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=11}
  %broadcast.120 = f32[10000,10]{1,0} broadcast(f32[10000]{0} %reshape.119), dimensions={0}, metadata={op_type="Mul" op_name="gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=11}
  %broadcast.100 = f32[10000,10]{1,0} broadcast(f32[10000]{0} %convert.80), dimensions={0}, metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %divide.101 = f32[10000,10]{1,0} divide(f32[10000,10]{1,0} %exponential.72, f32[10000,10]{1,0} %broadcast.100), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %subtract.102 = f32[10000,10]{1,0} subtract(f32[10000,10]{1,0} %divide.101, f32[10000,10]{1,0} %add.63), metadata={op_type="SparseSoftmaxCrossEntropyWithLogits" op_name="SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=7}
  %multiply.121 = f32[10000,10]{1,0} multiply(f32[10000,10]{1,0} %broadcast.120, f32[10000,10]{1,0} %subtract.102), metadata={op_type="Mul" op_name="gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=11}
  %dot.153 = f32[784,10]{1,0} dot(f32[10000,784]{1,0} %reshape.37, f32[10000,10]{1,0} %multiply.121), lhs_contracting_dims={0}, rhs_contracting_dims={0}, metadata={op_type="MatMul" op_name="gradient_tape/dense/MatMul" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=11}
  %transpose.154 = f32[784,10]{1,0} transpose(f32[784,10]{1,0} %dot.153), dimensions={0,1}, metadata={op_type="MatMul" op_name="gradient_tape/dense/MatMul" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=11}
  %subtract.163 = f32[784,10]{1,0} subtract(f32[784,10]{1,0} %transpose.154, f32[784,10]{1,0} %arg8.9), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %constant.156 = f32[] constant(1), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.162 = f32[] subtract(f32[] %constant.156, f32[] %arg6.7), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %broadcast.164 = f32[784,10]{1,0} broadcast(f32[] %subtract.162), dimensions={}, metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %multiply.165 = f32[784,10]{1,0} multiply(f32[784,10]{1,0} %subtract.163, f32[784,10]{1,0} %broadcast.164), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %add.166 = f32[784,10]{1,0} add(f32[784,10]{1,0} %arg8.9, f32[784,10]{1,0} %multiply.165), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.158 = f32[] subtract(f32[] %constant.156, f32[] %power.29), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %sqrt.159 = f32[] sqrt(f32[] %subtract.158), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %multiply.160 = f32[] multiply(f32[] %arg4.5, f32[] %sqrt.159), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.157 = f32[] subtract(f32[] %constant.156, f32[] %power.25), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %divide.161 = f32[] divide(f32[] %multiply.160, f32[] %subtract.157), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %broadcast.176 = f32[784,10]{1,0} broadcast(f32[] %divide.161), dimensions={}, metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %multiply.177 = f32[784,10]{1,0} multiply(f32[784,10]{1,0} %add.166, f32[784,10]{1,0} %broadcast.176), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %arg9.10 = f32[784,10]{1,0} parameter(9), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %multiply.168 = f32[784,10]{1,0} multiply(f32[784,10]{1,0} %transpose.154, f32[784,10]{1,0} %transpose.154), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.169 = f32[784,10]{1,0} subtract(f32[784,10]{1,0} %multiply.168, f32[784,10]{1,0} %arg9.10), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.167 = f32[] subtract(f32[] %constant.156, f32[] %arg7.8), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %broadcast.170 = f32[784,10]{1,0} broadcast(f32[] %subtract.167), dimensions={}, metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %multiply.171 = f32[784,10]{1,0} multiply(f32[784,10]{1,0} %subtract.169, f32[784,10]{1,0} %broadcast.170), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %add.172 = f32[784,10]{1,0} add(f32[784,10]{1,0} %arg9.10, f32[784,10]{1,0} %multiply.171), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %sqrt.173 = f32[784,10]{1,0} sqrt(f32[784,10]{1,0} %add.172), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %constant.155 = f32[] constant(1e-07), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %broadcast.174 = f32[784,10]{1,0} broadcast(f32[] %constant.155), dimensions={}, metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %add.175 = f32[784,10]{1,0} add(f32[784,10]{1,0} %sqrt.173, f32[784,10]{1,0} %broadcast.174), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %divide.178 = f32[784,10]{1,0} divide(f32[784,10]{1,0} %multiply.177, f32[784,10]{1,0} %add.175), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.179 = f32[784,10]{1,0} subtract(f32[784,10]{1,0} %arg2.3, f32[784,10]{1,0} %divide.178), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %reshape.182 = f32[784,10]{1,0} reshape(f32[784,10]{1,0} %subtract.179), metadata={op_name="XLA_Retvals"}
  %tuple.183 = (f32[784,10]{1,0}) tuple(f32[784,10]{1,0} %reshape.182), metadata={op_name="XLA_Retvals"}
  %get-tuple-element.184 = f32[784,10]{1,0} get-tuple-element((f32[784,10]{1,0}) %tuple.183), index=0, metadata={op_name="XLA_Retvals"}
  %arg10.11 = f32[10]{0} parameter(10), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %constant.122 = f32[] constant(0), metadata={op_type="BiasAddGrad" op_name="gradient_tape/dense/BiasAdd/BiasAddGrad" source_file="/tmp/ipykernel_31814/3922067182.py" source_line=11}
  %reduce.127 = f32[10]{0} reduce(f32[10000,10]{1,0} %multiply.121, f32[] %constant.122), dimensions={0}, to_apply=%region_0.123, metadata={op_name="gradient_tape/dense/BiasAdd/BiasAddGrad"}
  %subtract.136 = f32[10]{0} subtract(f32[10]{0} %reduce.127, f32[10]{0} %arg10.11), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %constant.129 = f32[] constant(1), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.135 = f32[] subtract(f32[] %constant.129, f32[] %arg6.7), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %broadcast.137 = f32[10]{0} broadcast(f32[] %subtract.135), dimensions={}, metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %multiply.138 = f32[10]{0} multiply(f32[10]{0} %subtract.136, f32[10]{0} %broadcast.137), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %add.139 = f32[10]{0} add(f32[10]{0} %arg10.11, f32[10]{0} %multiply.138), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.131 = f32[] subtract(f32[] %constant.129, f32[] %power.29), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %sqrt.132 = f32[] sqrt(f32[] %subtract.131), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %multiply.133 = f32[] multiply(f32[] %arg4.5, f32[] %sqrt.132), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.130 = f32[] subtract(f32[] %constant.129, f32[] %power.25), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %divide.134 = f32[] divide(f32[] %multiply.133, f32[] %subtract.130), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %broadcast.149 = f32[10]{0} broadcast(f32[] %divide.134), dimensions={}, metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %multiply.150 = f32[10]{0} multiply(f32[10]{0} %add.139, f32[10]{0} %broadcast.149), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %arg11.12 = f32[10]{0} parameter(11), parameter_replication={false}, metadata={op_name="XLA_Args"}
  %multiply.141 = f32[10]{0} multiply(f32[10]{0} %reduce.127, f32[10]{0} %reduce.127), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.142 = f32[10]{0} subtract(f32[10]{0} %multiply.141, f32[10]{0} %arg11.12), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.140 = f32[] subtract(f32[] %constant.129, f32[] %arg7.8), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %broadcast.143 = f32[10]{0} broadcast(f32[] %subtract.140), dimensions={}, metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %multiply.144 = f32[10]{0} multiply(f32[10]{0} %subtract.142, f32[10]{0} %broadcast.143), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %add.145 = f32[10]{0} add(f32[10]{0} %arg11.12, f32[10]{0} %multiply.144), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %sqrt.146 = f32[10]{0} sqrt(f32[10]{0} %add.145), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %constant.128 = f32[] constant(1e-07), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %broadcast.147 = f32[10]{0} broadcast(f32[] %constant.128), dimensions={}, metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %add.148 = f32[10]{0} add(f32[10]{0} %sqrt.146, f32[10]{0} %broadcast.147), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %divide.151 = f32[10]{0} divide(f32[10]{0} %multiply.150, f32[10]{0} %add.148), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %subtract.152 = f32[10]{0} subtract(f32[10]{0} %arg3.4, f32[10]{0} %divide.151), metadata={op_type="ResourceApplyAdam" op_name="Adam/Adam/update_1/ResourceApplyAdam" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/adam.py" source_line=176}
  %reshape.185 = f32[10]{0} reshape(f32[10]{0} %subtract.152), metadata={op_name="XLA_Retvals"}
  %tuple.186 = (f32[10]{0}) tuple(f32[10]{0} %reshape.185), metadata={op_name="XLA_Retvals"}
  %get-tuple-element.187 = f32[10]{0} get-tuple-element((f32[10]{0}) %tuple.186), index=0, metadata={op_name="XLA_Retvals"}
  %constant.180 = s64[] constant(1), metadata={op_type="AssignAddVariableOp" op_name="Adam/Adam/AssignAddVariableOp" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py" source_line=725}
  %add.181 = s64[] add(s64[] %arg5.6, s64[] %constant.180), metadata={op_type="AssignAddVariableOp" op_name="Adam/Adam/AssignAddVariableOp" source_file="/tmpfs/src/tf_docs_env/lib/python3.7/site-packages/keras/optimizer_v2/optimizer_v2.py" source_line=725}
  %reshape.188 = s64[] reshape(s64[] %add.181), metadata={op_name="XLA_Retvals"}
  %tuple.189 = (s64[]) tuple(s64[] %reshape.188), metadata={op_name="XLA_Retvals"}
  %get-tuple-element.190 = s64[] get-tuple-element((s64[]) %tuple.189), index=0, metadata={op_name="XLA_Retvals"}
  %reshape.191 = f32[784,10]{1,0} reshape(f32[784,10]{1,0} %add.166), metadata={op_name="XLA_Retvals"}
  %tuple.192 = (f32[784,10]{1,0}) tuple(f32[784,10]{1,0} %reshape.191), metadata={op_name="XLA_Retvals"}
  %get-tuple-element.193 = f32[784,10]{1,0} get-tuple-element((f32[784,10]{1,0}) %tuple.192), index=0, metadata={op_name="XLA_Retvals"}
  %reshape.194 = f32[784,10]{1,0} reshape(f32[784,10]{1,0} %add.172), metadata={op_name="XLA_Retvals"}
  %tuple.195 = (f32[784,10]{1,0}) tuple(f32[784,10]{1,0} %reshape.194), metadata={op_name="XLA_Retvals"}
  %get-tuple-element.196 = f32[784,10]{1,0} get-tuple-element((f32[784,10]{1,0}) %tuple.195), index=0, metadata={op_name="XLA_Retvals"}
  %reshape.197 = f32[10]{0} reshape(f32[10]{0} %add.139), metadata={op_name="XLA_Retvals"}
  %tuple.198 = (f32[10]{0}) tuple(f32[10]{0} %reshape.197), metadata={op_name="XLA_Retvals"}
  %get-tuple-element.199 = f32[10]{0} get-tuple-element((f32[10]{0}) %tuple.198), index=0, metadata={op_name="XLA_Retvals"}
  %reshape.200 = f32[10]{0} reshape(f32[10]{0} %add.145), metadata={op_name="XLA_Retvals"}
  %tuple.201 = (f32[10]{0}) tuple(f32[10]{0} %reshape.200), metadata={op_name="XLA_Retvals"}
  %get-tuple-element.202 = f32[10]{0} get-tuple-element((f32[10]{0}) %tuple.201), index=0, metadata={op_name="XLA_Retvals"}
  ROOT %tuple.203 = (f32[784,10]{1,0}, f32[10]{0}, s64[], f32[784,10]{1,0}, f32[784,10]{1,0}, /*index=5*/f32[10]{0}, f32[10]{0}) tuple(f32[784,10]{1,0} %get-tuple-element.184, f32[10]{0} %get-tuple-element.187, s64[] %get-tuple-element.190, f32[784,10]{1,0} %get-tuple-element.193, f32[784,10]{1,0} %get-tuple-element.196, /*index=5*/f32[10]{0} %get-tuple-element.199, f32[10]{0} %get-tuple-element.202), metadata={op_name="XLA_Retvals"}
}