{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "f4TSNCvpENrW" }, "source": [ "##### Copyright 2019 The TensorFlow Authors." ] }, { "cell_type": "code", "execution_count": 1, "metadata": { "cellView": "form", "execution": { "iopub.execute_input": "2023-01-18T12:09:52.208250Z", "iopub.status.busy": "2023-01-18T12:09:52.208024Z", "iopub.status.idle": "2023-01-18T12:09:52.211701Z", "shell.execute_reply": "2023-01-18T12:09:52.211136Z" }, "id": "vamNSA0vEP-m" }, "outputs": [], "source": [ "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n", "# you may not use this file except in compliance with the License.\n", "# You may obtain a copy of the License at\n", "#\n", "# https://www.apache.org/licenses/LICENSE-2.0\n", "#\n", "# Unless required by applicable law or agreed to in writing, software\n", "# distributed under the License is distributed on an \"AS IS\" BASIS,\n", "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n", "# See the License for the specific language governing permissions and\n", "# limitations under the License." ] }, { "cell_type": "markdown", "metadata": { "id": "e1oSi4lHFt3z" }, "source": [ "# Use XLA with tf.function" ] }, { "cell_type": "markdown", "metadata": { "id": "b7noD9NjFRL-" }, "source": [ "\n", " \n", " \n", " \n", " \n", "
\n", " View on TensorFlow.org\n", " \n", " Run in Google Colab\n", " \n", " Download notebook\n", " \n", " View source on GitHub\n", "
" ] }, { "cell_type": "markdown", "metadata": { "id": "sDy5lSBd4BDE" }, "source": [ "This tutorial trains a TensorFlow model to classify the MNIST dataset, where the training function is compiled using XLA.\n", "\n", "First, load TensorFlow and enable eager execution." ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "execution": { "iopub.execute_input": "2023-01-18T12:09:52.215086Z", "iopub.status.busy": "2023-01-18T12:09:52.214569Z", "iopub.status.idle": "2023-01-18T12:09:54.132611Z", "shell.execute_reply": "2023-01-18T12:09:54.131920Z" }, "id": "45kUPj5ZFrRa" }, "outputs": [ { "name": "stderr", "output_type": "stream", "text": [ "2023-01-18 12:09:53.148023: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory\n", "2023-01-18 12:09:53.148129: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory\n", "2023-01-18 12:09:53.148139: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.\n" ] } ], "source": [ "import tensorflow as tf\n" ] }, { "cell_type": "markdown", "metadata": { "id": "GZVNiRmTDV-5" }, "source": [ "Then define some necessary constants and prepare the MNIST dataset." ] }, { "cell_type": "code", "execution_count": 3, "metadata": { "execution": { "iopub.execute_input": "2023-01-18T12:09:54.136770Z", "iopub.status.busy": "2023-01-18T12:09:54.136143Z", "iopub.status.idle": "2023-01-18T12:09:58.448922Z", "shell.execute_reply": "2023-01-18T12:09:58.448213Z" }, "id": "f37TSEGvGX4_" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz\n" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\r", " 8192/11490434 [..............................] - ETA: 0s" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r", " 6971392/11490434 [=================>............] - ETA: 0s" ] }, { "name": "stdout", "output_type": "stream", "text": [ "\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\b\r", "11490434/11490434 [==============================] - 0s 0us/step\n" ] } ], "source": [ "# Size of each input image, 28 x 28 pixels\n", "IMAGE_SIZE = 28 * 28\n", "# Number of distinct number labels, [0..9]\n", "NUM_CLASSES = 10\n", "# Number of examples in each training batch (step)\n", "TRAIN_BATCH_SIZE = 100\n", "# Number of training steps to run\n", "TRAIN_STEPS = 1000\n", "\n", "# Loads MNIST dataset.\n", "train, test = tf.keras.datasets.mnist.load_data()\n", "train_ds = tf.data.Dataset.from_tensor_slices(train).batch(TRAIN_BATCH_SIZE).repeat()\n", "\n", "# Casting from raw data to the required datatypes.\n", "def cast(images, labels):\n", " images = tf.cast(\n", " tf.reshape(images, [-1, IMAGE_SIZE]), tf.float32)\n", " labels = tf.cast(labels, tf.int64)\n", " return (images, labels)" ] }, { "cell_type": "markdown", "metadata": { "id": "lv7I-u_82v1S" }, "source": [ "Finally, define the model and the optimizer. The model uses a single dense layer." ] }, { "cell_type": "code", "execution_count": 4, "metadata": { "execution": { "iopub.execute_input": "2023-01-18T12:09:58.453106Z", "iopub.status.busy": "2023-01-18T12:09:58.452620Z", "iopub.status.idle": "2023-01-18T12:09:58.466177Z", "shell.execute_reply": "2023-01-18T12:09:58.465533Z" }, "id": "7O2NcEfG206Q" }, "outputs": [], "source": [ "layer = tf.keras.layers.Dense(NUM_CLASSES)\n", "optimizer = tf.keras.optimizers.Adam()" ] }, { "cell_type": "markdown", "metadata": { "id": "x_ZehpZP-SfS" }, "source": [ "# Define the training function\n", "\n", "In the training function, you get the predicted labels using the layer defined above, and then minimize the gradient of the loss using the optimizer. In order to compile the computation using XLA, place it inside `tf.function` with `jit_compile=True`." ] }, { "cell_type": "code", "execution_count": 5, "metadata": { "execution": { "iopub.execute_input": "2023-01-18T12:09:58.469693Z", "iopub.status.busy": "2023-01-18T12:09:58.469106Z", "iopub.status.idle": "2023-01-18T12:09:58.473553Z", "shell.execute_reply": "2023-01-18T12:09:58.472998Z" }, "id": "ZbhJl_WvGa3g" }, "outputs": [], "source": [ "@tf.function(jit_compile=True)\n", "def train_mnist(images, labels):\n", " images, labels = cast(images, labels)\n", "\n", " with tf.GradientTape() as tape:\n", " predicted_labels = layer(images)\n", " loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(\n", " logits=predicted_labels, labels=labels\n", " ))\n", " layer_variables = layer.trainable_variables\n", " grads = tape.gradient(loss, layer_variables)\n", " optimizer.apply_gradients(zip(grads, layer_variables))" ] }, { "cell_type": "markdown", "metadata": { "id": "EZD1m_n1DxAF" }, "source": [ "# Train and test the model" ] }, { "cell_type": "markdown", "metadata": { "id": "gukC2Hol3sFZ" }, "source": [ "Once you have defined the training function, define the model." ] }, { "cell_type": "code", "execution_count": 6, "metadata": { "execution": { "iopub.execute_input": "2023-01-18T12:09:58.476699Z", "iopub.status.busy": "2023-01-18T12:09:58.476250Z", "iopub.status.idle": "2023-01-18T12:10:02.271277Z", "shell.execute_reply": "2023-01-18T12:10:02.270515Z" }, "id": "qe28bAHNHUG2" }, "outputs": [], "source": [ "for images, labels in train_ds:\n", " if optimizer.iterations > TRAIN_STEPS:\n", " break\n", " train_mnist(images, labels)" ] }, { "cell_type": "markdown", "metadata": { "id": "qgsKmz3n2UiW" }, "source": [ "And, finally, check the accuracy:" ] }, { "cell_type": "code", "execution_count": 7, "metadata": { "execution": { "iopub.execute_input": "2023-01-18T12:10:02.275609Z", "iopub.status.busy": "2023-01-18T12:10:02.275017Z", "iopub.status.idle": "2023-01-18T12:10:02.301542Z", "shell.execute_reply": "2023-01-18T12:10:02.300860Z" }, "id": "_GxF6jTRHVuA" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Prediction accuracy after training: tf.Tensor(0.8773, shape=(), dtype=float32)\n" ] } ], "source": [ "images, labels = cast(test[0], test[1])\n", "predicted_labels = layer(images)\n", "correct_prediction = tf.equal(tf.argmax(predicted_labels, 1), labels)\n", "accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))\n", "print(\"Prediction accuracy after training: %s\" % accuracy)" ] }, { "cell_type": "markdown", "metadata": { "id": "PXoOjJnuZRaV" }, "source": [ "Behind the scenes, the XLA compiler has compiled the entire TF function to HLO, which has enabled fusion optimizations. Using the introspection facilities, we can see the HLO code (other interesting possible values for \"stage\" are `optimized_hlo` for HLO after optimizations and `optimized_hlo_dot` for a Graphviz graph):" ] }, { "cell_type": "code", "execution_count": 8, "metadata": { "execution": { "iopub.execute_input": "2023-01-18T12:10:02.304801Z", "iopub.status.busy": "2023-01-18T12:10:02.304322Z", "iopub.status.idle": "2023-01-18T12:10:02.369467Z", "shell.execute_reply": "2023-01-18T12:10:02.368800Z" }, "id": "_a8GsNLVaLSQ" }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "HloModule a_inference_train_mnist_5324__.206, input_output_alias={ {0}: (2, {}, may-alias), {1}: (3, {}, may-alias), {2}: (4, {}, may-alias), {3}: (6, {}, may-alias), {4}: (7, {}, may-alias), {5}: (8, {}, may-alias), {6}: (9, {}, may-alias) }, entry_computation_layout={(f32[10000,784]{1,0},s64[10000]{0},f32[784,10]{1,0},f32[10]{0},s64[],f32[],f32[784,10]{1,0},f32[784,10]{1,0},f32[10]{0},f32[10]{0})->(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})}\n", "\n", "%max_float_.42 (x.43: f32[], y.44: f32[]) -> f32[] {\n", " %x.43 = f32[] parameter(0)\n", " %y.44 = f32[] parameter(1)\n", " ROOT %maximum.45 = f32[] maximum(f32[] %x.43, f32[] %y.44)\n", "}\n", "\n", "%add_float_.52 (x.53: f32[], y.54: f32[]) -> f32[] {\n", " %x.53 = f32[] parameter(0)\n", " %y.54 = f32[] parameter(1)\n", " ROOT %add.55 = f32[] add(f32[] %x.53, f32[] %y.54)\n", "}\n", "\n", "%add_float_.71 (x.72: f32[], y.73: f32[]) -> f32[] {\n", " %x.72 = f32[] parameter(0)\n", " %y.73 = f32[] parameter(1)\n", " ROOT %add.74 = f32[] add(f32[] %x.72, f32[] %y.73)\n", "}\n", "\n", "%Mean-reduction.83 (x.84: f32[], y.85: f32[]) -> f32[] {\n", " %x.84 = f32[] parameter(0)\n", " %y.85 = f32[] parameter(1)\n", " ROOT %add.86 = f32[] add(f32[] %x.84, f32[] %y.85)\n", "}\n", "\n", "%region_0.100 (Arg_0.101: f32[], Arg_1.102: f32[]) -> f32[] {\n", " %Arg_0.101 = f32[] parameter(0)\n", " %Arg_1.102 = f32[] parameter(1)\n", " ROOT %add.103 = f32[] add(f32[] %Arg_0.101, f32[] %Arg_1.102), metadata={op_type=\"BiasAddGrad\" op_name=\"gradient_tape/dense/BiasAdd/BiasAddGrad\"}\n", "}\n", "\n", "ENTRY %a_inference_train_mnist_5324__.206 (arg0.1: f32[10000,784], arg1.2: s64[10000], arg2.3: f32[784,10], arg3.4: f32[10], arg4.5: s64[], arg5.6: f32[], arg6.7: f32[784,10], arg7.8: f32[784,10], arg8.9: f32[10], arg9.10: f32[10]) -> (f32[784,10], f32[10], s64[], f32[784,10], f32[784,10], /*index=5*/f32[10], f32[10]) {\n", " %constant.13 = s32[2]{0} constant({-1, 784}), metadata={op_type=\"Reshape\" op_name=\"Reshape\" source_file=\"/tmpfs/tmp/ipykernel_11149/494562224.py\" source_line=16}\n", " %arg1.2 = s64[10000]{0} parameter(1), parameter_replication={false}, metadata={op_name=\"XLA_Args\"}\n", " %reshape.12 = s64[10000]{0} reshape(s64[10000]{0} %arg1.2)\n", " %broadcast.24 = s64[10000,10]{1,0} broadcast(s64[10000]{0} %reshape.12), dimensions={0}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %iota.23 = s64[10000,10]{1,0} iota(), iota_dimension=1, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %compare.25 = pred[10000,10]{1,0} compare(s64[10000,10]{1,0} %broadcast.24, s64[10000,10]{1,0} %iota.23), direction=EQ, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.20 = f32[] constant(1), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.22 = f32[10000,10]{1,0} broadcast(f32[] %constant.20), dimensions={}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.19 = f32[] constant(0), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.21 = f32[10000,10]{1,0} broadcast(f32[] %constant.19), dimensions={}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %select.26 = f32[10000,10]{1,0} select(pred[10000,10]{1,0} %compare.25, f32[10000,10]{1,0} %broadcast.22, f32[10000,10]{1,0} %broadcast.21), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.34 = s64[] constant(0), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.35 = s64[10000]{0} broadcast(s64[] %constant.34), dimensions={}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %compare.36 = pred[10000]{0} compare(s64[10000]{0} %broadcast.35, s64[10000]{0} %reshape.12), direction=LE, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.31 = s64[] constant(10), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.32 = s64[10000]{0} broadcast(s64[] %constant.31), dimensions={}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %compare.33 = pred[10000]{0} compare(s64[10000]{0} %reshape.12, s64[10000]{0} %broadcast.32), direction=LT, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %and.37 = pred[10000]{0} and(pred[10000]{0} %compare.36, pred[10000]{0} %compare.33), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.29 = f32[] constant(0), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.30 = f32[10000]{0} broadcast(f32[] %constant.29), dimensions={}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.27 = f32[] constant(nan), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.28 = f32[10000]{0} broadcast(f32[] %constant.27), dimensions={}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %select.38 = f32[10000]{0} select(pred[10000]{0} %and.37, f32[10000]{0} %broadcast.30, f32[10000]{0} %broadcast.28), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.39 = f32[10000,10]{1,0} broadcast(f32[10000]{0} %select.38), dimensions={0}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %add.40 = f32[10000,10]{1,0} add(f32[10000,10]{1,0} %select.26, f32[10000,10]{1,0} %broadcast.39), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %negate.67 = f32[10000,10]{1,0} negate(f32[10000,10]{1,0} %add.40), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.63 = f32[] constant(0), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.64 = f32[10000,10]{1,0} broadcast(f32[] %constant.63), dimensions={}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %compare.65 = pred[10000,10]{1,0} compare(f32[10000,10]{1,0} %add.40, f32[10000,10]{1,0} %broadcast.64), direction=EQ, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.61 = f32[] constant(0), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.62 = f32[10000,10]{1,0} broadcast(f32[] %constant.61), dimensions={}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %arg0.1 = f32[10000,784]{1,0} parameter(0), parameter_replication={false}, metadata={op_name=\"XLA_Args\"}\n", " %reshape.11 = f32[10000,784]{1,0} reshape(f32[10000,784]{1,0} %arg0.1)\n", " %reshape.14 = f32[10000,784]{1,0} reshape(f32[10000,784]{1,0} %reshape.11), metadata={op_type=\"Reshape\" op_name=\"Reshape\" source_file=\"/tmpfs/tmp/ipykernel_11149/494562224.py\" source_line=16}\n", " %arg2.3 = f32[784,10]{1,0} parameter(2), parameter_replication={false}, metadata={op_name=\"XLA_Args\"}\n", " %dot.15 = f32[10000,10]{1,0} dot(f32[10000,784]{1,0} %reshape.14, 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.9/site-packages/keras/layers/core/dense.py\" source_line=241}\n", " %transpose.16 = f32[10000,10]{1,0} transpose(f32[10000,10]{1,0} %dot.15), dimensions={0,1}, metadata={op_type=\"MatMul\" op_name=\"dense/MatMul\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/layers/core/dense.py\" source_line=241}\n", " %arg3.4 = f32[10]{0} parameter(3), parameter_replication={false}, metadata={op_name=\"XLA_Args\"}\n", " %broadcast.17 = 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.9/site-packages/keras/layers/core/dense.py\" source_line=252}\n", " %add.18 = f32[10000,10]{1,0} add(f32[10000,10]{1,0} %transpose.16, f32[10000,10]{1,0} %broadcast.17), metadata={op_type=\"BiasAdd\" op_name=\"dense/BiasAdd\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/layers/core/dense.py\" source_line=252}\n", " %constant.41 = f32[] constant(-inf), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %reduce.46 = f32[10000]{0} reduce(f32[10000,10]{1,0} %add.18, f32[] %constant.41), dimensions={1}, to_apply=%max_float_.42, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.47 = f32[10000,10]{1,0} broadcast(f32[10000]{0} %reduce.46), dimensions={0}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %subtract.48 = f32[10000,10]{1,0} subtract(f32[10000,10]{1,0} %add.18, f32[10000,10]{1,0} %broadcast.47), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %exponential.49 = f32[10000,10]{1,0} exponential(f32[10000,10]{1,0} %subtract.48), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %convert.50 = f32[10000,10]{1,0} convert(f32[10000,10]{1,0} %exponential.49), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.51 = f32[] constant(0), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %reduce.56 = f32[10000]{0} reduce(f32[10000,10]{1,0} %convert.50, f32[] %constant.51), dimensions={1}, to_apply=%add_float_.52, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %convert.57 = f32[10000]{0} convert(f32[10000]{0} %reduce.56), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %log.58 = f32[10000]{0} log(f32[10000]{0} %convert.57), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %broadcast.59 = f32[10000,10]{1,0} broadcast(f32[10000]{0} %log.58), dimensions={0}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %subtract.60 = f32[10000,10]{1,0} subtract(f32[10000,10]{1,0} %subtract.48, f32[10000,10]{1,0} %broadcast.59), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %select.66 = f32[10000,10]{1,0} select(pred[10000,10]{1,0} %compare.65, f32[10000,10]{1,0} %broadcast.62, f32[10000,10]{1,0} %subtract.60), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %multiply.68 = f32[10000,10]{1,0} multiply(f32[10000,10]{1,0} %negate.67, f32[10000,10]{1,0} %select.66), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %convert.70 = f32[10000,10]{1,0} convert(f32[10000,10]{1,0} %multiply.68), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.69 = f32[] constant(0), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %reduce.75 = f32[10000]{0} reduce(f32[10000,10]{1,0} %convert.70, f32[] %constant.69), dimensions={1}, to_apply=%add_float_.71, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %convert.76 = f32[10000]{0} convert(f32[10000]{0} %reduce.75), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %convert.80 = f32[10000]{0} convert(f32[10000]{0} %convert.76), metadata={op_type=\"Mean\" op_name=\"Mean\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.81 = f32[] constant(0), metadata={op_type=\"Mean\" op_name=\"Mean\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %convert.82 = f32[] convert(f32[] %constant.81), metadata={op_type=\"Mean\" op_name=\"Mean\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %reduce.87 = f32[] reduce(f32[10000]{0} %convert.80, f32[] %convert.82), dimensions={0}, to_apply=%Mean-reduction.83, metadata={op_type=\"Mean\" op_name=\"Mean\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.88 = s32[] constant(10000), metadata={op_type=\"Mean\" op_name=\"Mean\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %convert.89 = f32[] convert(s32[] %constant.88), metadata={op_type=\"Mean\" op_name=\"Mean\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %divide.90 = f32[] divide(f32[] %reduce.87, f32[] %convert.89), metadata={op_type=\"Mean\" op_name=\"Mean\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %convert.91 = f32[] convert(f32[] %divide.90), metadata={op_type=\"Mean\" op_name=\"Mean\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %constant.92 = f32[] constant(0.0001), metadata={op_type=\"Mul\" op_name=\"gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=11}\n", " %broadcast.93 = f32[10000,1]{1,0} broadcast(f32[] %constant.92), dimensions={}, metadata={op_type=\"Mul\" op_name=\"gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=11}\n", " %constant.109 = s64[] constant(1), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall/add\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=162}\n", " %constant.113 = f32[] constant(0.9), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall/Pow\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=163}\n", " %constant.116 = f32[] constant(1), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall/sub_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %constant.119 = f32[] constant(0.999), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall/Pow_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=164}\n", " %constant.122 = f32[] constant(1), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall/sub\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %constant.126 = s64[] constant(1), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall_1/add\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=162}\n", " %constant.130 = f32[] constant(0.9), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall_1/Pow\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=163}\n", " %constant.133 = f32[] constant(1), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall_1/sub_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %constant.136 = f32[] constant(0.999), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall_1/Pow_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=164}\n", " %constant.139 = f32[] constant(1), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall_1/sub\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %constant.148 = f32[] constant(0.1), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %constant.156 = f32[] constant(0.001), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul_2\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %constant.162 = f32[] constant(1e-07), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall/add_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %constant.169 = f32[] constant(0.1), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %constant.177 = f32[] constant(0.001), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul_2\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %constant.183 = f32[] constant(1e-07), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall_1/add_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %arg6.7 = f32[784,10]{1,0} parameter(6), parameter_replication={false}, metadata={op_name=\"XLA_Args\"}\n", " %constant.94 = f32[] constant(0.0001), metadata={op_type=\"Mul\" op_name=\"gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=11}\n", " %broadcast.95 = f32[10000,1]{1,0} broadcast(f32[] %constant.94), dimensions={}, metadata={op_type=\"Mul\" op_name=\"gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=11}\n", " %reshape.96 = f32[10000]{0} reshape(f32[10000,1]{1,0} %broadcast.95), metadata={op_type=\"Mul\" op_name=\"gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=11}\n", " %broadcast.97 = f32[10000,10]{1,0} broadcast(f32[10000]{0} %reshape.96), dimensions={0}, metadata={op_type=\"Mul\" op_name=\"gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=11}\n", " %broadcast.77 = f32[10000,10]{1,0} broadcast(f32[10000]{0} %convert.57), dimensions={0}, metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %divide.78 = f32[10000,10]{1,0} divide(f32[10000,10]{1,0} %exponential.49, f32[10000,10]{1,0} %broadcast.77), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %subtract.79 = f32[10000,10]{1,0} subtract(f32[10000,10]{1,0} %divide.78, f32[10000,10]{1,0} %add.40), metadata={op_type=\"SparseSoftmaxCrossEntropyWithLogits\" op_name=\"SparseSoftmaxCrossEntropyWithLogits/SparseSoftmaxCrossEntropyWithLogits\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=7}\n", " %multiply.98 = f32[10000,10]{1,0} multiply(f32[10000,10]{1,0} %broadcast.97, f32[10000,10]{1,0} %subtract.79), metadata={op_type=\"Mul\" op_name=\"gradient_tape/SparseSoftmaxCrossEntropyWithLogits/mul\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=11}\n", " %dot.105 = f32[784,10]{1,0} dot(f32[10000,784]{1,0} %reshape.14, f32[10000,10]{1,0} %multiply.98), lhs_contracting_dims={0}, rhs_contracting_dims={0}, metadata={op_type=\"MatMul\" op_name=\"gradient_tape/dense/MatMul/MatMul\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=11}\n", " %transpose.106 = f32[784,10]{1,0} transpose(f32[784,10]{1,0} %dot.105), dimensions={0,1}, metadata={op_type=\"MatMul\" op_name=\"gradient_tape/dense/MatMul/MatMul\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=11}\n", " %subtract.147 = f32[784,10]{1,0} subtract(f32[784,10]{1,0} %transpose.106, f32[784,10]{1,0} %arg6.7), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall/sub_2\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %constant.149 = f32[] constant(0.1), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %broadcast.150 = f32[784,10]{1,0} broadcast(f32[] %constant.149), dimensions={}, metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %multiply.151 = f32[784,10]{1,0} multiply(f32[784,10]{1,0} %subtract.147, f32[784,10]{1,0} %broadcast.150), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %add.152 = f32[784,10]{1,0} add(f32[784,10]{1,0} %arg6.7, f32[784,10]{1,0} %multiply.151), metadata={op_type=\"AssignAddVariableOp\" op_name=\"StatefulPartitionedCall/AssignAddVariableOp\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %arg5.6 = f32[] parameter(5), parameter_replication={false}, metadata={op_name=\"XLA_Args\"}\n", " %constant.123 = f32[] constant(1), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall/sub\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %constant.120 = f32[] constant(0.999), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall/Pow_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=164}\n", " %arg4.5 = s64[] parameter(4), parameter_replication={false}, metadata={op_name=\"XLA_Args\"}\n", " %constant.110 = s64[] constant(1), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall/add\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=162}\n", " %add.111 = s64[] add(s64[] %arg4.5, s64[] %constant.110), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall/add\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=162}\n", " %convert.112 = f32[] convert(s64[] %add.111), metadata={op_type=\"Cast\" op_name=\"StatefulPartitionedCall/Cast\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=162}\n", " %power.121 = f32[] power(f32[] %constant.120, f32[] %convert.112), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall/Pow_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=164}\n", " %subtract.124 = f32[] subtract(f32[] %constant.123, f32[] %power.121), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall/sub\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %sqrt.125 = f32[] sqrt(f32[] %subtract.124), metadata={op_type=\"Sqrt\" op_name=\"Sqrt.StatefulPartitionedCall/Sqrt\"}\n", " %multiply.143 = f32[] multiply(f32[] %arg5.6, f32[] %sqrt.125), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %constant.117 = f32[] constant(1), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall/sub_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %constant.114 = f32[] constant(0.9), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall/Pow\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=163}\n", " %power.115 = f32[] power(f32[] %constant.114, f32[] %convert.112), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall/Pow\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=163}\n", " %subtract.118 = f32[] subtract(f32[] %constant.117, f32[] %power.115), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall/sub_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %divide.144 = f32[] divide(f32[] %multiply.143, f32[] %subtract.118), metadata={op_type=\"RealDiv\" op_name=\"StatefulPartitionedCall/truediv\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %broadcast.153 = f32[784,10]{1,0} broadcast(f32[] %divide.144), dimensions={}, metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul_3\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %multiply.154 = f32[784,10]{1,0} multiply(f32[784,10]{1,0} %add.152, f32[784,10]{1,0} %broadcast.153), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul_3\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %arg7.8 = f32[784,10]{1,0} parameter(7), parameter_replication={false}, metadata={op_name=\"XLA_Args\"}\n", " %multiply.107 = f32[784,10]{1,0} multiply(f32[784,10]{1,0} %transpose.106, f32[784,10]{1,0} %transpose.106), metadata={op_type=\"Square\" op_name=\"StatefulPartitionedCall/Square\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %subtract.155 = f32[784,10]{1,0} subtract(f32[784,10]{1,0} %multiply.107, f32[784,10]{1,0} %arg7.8), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall/sub_3\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %constant.157 = f32[] constant(0.001), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul_2\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %broadcast.158 = f32[784,10]{1,0} broadcast(f32[] %constant.157), dimensions={}, metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul_2\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %multiply.159 = f32[784,10]{1,0} multiply(f32[784,10]{1,0} %subtract.155, f32[784,10]{1,0} %broadcast.158), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall/mul_2\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %add.160 = f32[784,10]{1,0} add(f32[784,10]{1,0} %arg7.8, f32[784,10]{1,0} %multiply.159), metadata={op_type=\"AssignAddVariableOp\" op_name=\"StatefulPartitionedCall/AssignAddVariableOp_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %sqrt.161 = f32[784,10]{1,0} sqrt(f32[784,10]{1,0} %add.160), metadata={op_type=\"Sqrt\" op_name=\"Sqrt_1.StatefulPartitionedCall/Sqrt_1\"}\n", " %constant.163 = f32[] constant(1e-07), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall/add_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %broadcast.164 = f32[784,10]{1,0} broadcast(f32[] %constant.163), dimensions={}, metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall/add_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %add.165 = f32[784,10]{1,0} add(f32[784,10]{1,0} %sqrt.161, f32[784,10]{1,0} %broadcast.164), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall/add_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %divide.166 = f32[784,10]{1,0} divide(f32[784,10]{1,0} %multiply.154, f32[784,10]{1,0} %add.165), metadata={op_type=\"RealDiv\" op_name=\"StatefulPartitionedCall/truediv_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %subtract.167 = f32[784,10]{1,0} subtract(f32[784,10]{1,0} %arg2.3, f32[784,10]{1,0} %divide.166), metadata={op_type=\"AssignSubVariableOp\" op_name=\"StatefulPartitionedCall/AssignSubVariableOp\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %reshape.191 = f32[784,10]{1,0} reshape(f32[784,10]{1,0} %subtract.167), metadata={op_name=\"XLA_Retvals\"}\n", " %copy.192 = f32[784,10]{1,0} copy(f32[784,10]{1,0} %reshape.191), metadata={op_name=\"XLA_Retvals\"}\n", " %arg8.9 = f32[10]{0} parameter(8), parameter_replication={false}, metadata={op_name=\"XLA_Args\"}\n", " %constant.99 = f32[] constant(-0), metadata={op_type=\"BiasAddGrad\" op_name=\"gradient_tape/dense/BiasAdd/BiasAddGrad\" source_file=\"/tmpfs/tmp/ipykernel_11149/3922067182.py\" source_line=11}\n", " %reduce.104 = f32[10]{0} reduce(f32[10000,10]{1,0} %multiply.98, f32[] %constant.99), dimensions={0}, to_apply=%region_0.100, metadata={op_type=\"BiasAddGrad\" op_name=\"gradient_tape/dense/BiasAdd/BiasAddGrad\"}\n", " %subtract.168 = f32[10]{0} subtract(f32[10]{0} %reduce.104, f32[10]{0} %arg8.9), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall_1/sub_2\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %constant.170 = f32[] constant(0.1), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %broadcast.171 = f32[10]{0} broadcast(f32[] %constant.170), dimensions={}, metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %multiply.172 = f32[10]{0} multiply(f32[10]{0} %subtract.168, f32[10]{0} %broadcast.171), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %add.173 = f32[10]{0} add(f32[10]{0} %arg8.9, f32[10]{0} %multiply.172), metadata={op_type=\"AssignAddVariableOp\" op_name=\"StatefulPartitionedCall_1/AssignAddVariableOp\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=194}\n", " %constant.140 = f32[] constant(1), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall_1/sub\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %constant.137 = f32[] constant(0.999), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall_1/Pow_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=164}\n", " %constant.127 = s64[] constant(1), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall_1/add\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=162}\n", " %add.128 = s64[] add(s64[] %arg4.5, s64[] %constant.127), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall_1/add\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=162}\n", " %convert.129 = f32[] convert(s64[] %add.128), metadata={op_type=\"Cast\" op_name=\"StatefulPartitionedCall_1/Cast\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=162}\n", " %power.138 = f32[] power(f32[] %constant.137, f32[] %convert.129), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall_1/Pow_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=164}\n", " %subtract.141 = f32[] subtract(f32[] %constant.140, f32[] %power.138), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall_1/sub\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %sqrt.142 = f32[] sqrt(f32[] %subtract.141), metadata={op_type=\"Sqrt\" op_name=\"Sqrt.StatefulPartitionedCall_1/Sqrt\"}\n", " %multiply.145 = f32[] multiply(f32[] %arg5.6, f32[] %sqrt.142), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %constant.134 = f32[] constant(1), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall_1/sub_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %constant.131 = f32[] constant(0.9), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall_1/Pow\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=163}\n", " %power.132 = f32[] power(f32[] %constant.131, f32[] %convert.129), metadata={op_type=\"Pow\" op_name=\"StatefulPartitionedCall_1/Pow\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=163}\n", " %subtract.135 = f32[] subtract(f32[] %constant.134, f32[] %power.132), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall_1/sub_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %divide.146 = f32[] divide(f32[] %multiply.145, f32[] %subtract.135), metadata={op_type=\"RealDiv\" op_name=\"StatefulPartitionedCall_1/truediv\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=170}\n", " %broadcast.174 = f32[10]{0} broadcast(f32[] %divide.146), dimensions={}, metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul_3\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %multiply.175 = f32[10]{0} multiply(f32[10]{0} %add.173, f32[10]{0} %broadcast.174), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul_3\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %arg9.10 = f32[10]{0} parameter(9), parameter_replication={false}, metadata={op_name=\"XLA_Args\"}\n", " %multiply.108 = f32[10]{0} multiply(f32[10]{0} %reduce.104, f32[10]{0} %reduce.104), metadata={op_type=\"Square\" op_name=\"StatefulPartitionedCall_1/Square\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %subtract.176 = f32[10]{0} subtract(f32[10]{0} %multiply.108, f32[10]{0} %arg9.10), metadata={op_type=\"Sub\" op_name=\"StatefulPartitionedCall_1/sub_3\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %constant.178 = f32[] constant(0.001), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul_2\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %broadcast.179 = f32[10]{0} broadcast(f32[] %constant.178), dimensions={}, metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul_2\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %multiply.180 = f32[10]{0} multiply(f32[10]{0} %subtract.176, f32[10]{0} %broadcast.179), metadata={op_type=\"Mul\" op_name=\"StatefulPartitionedCall_1/mul_2\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %add.181 = f32[10]{0} add(f32[10]{0} %arg9.10, f32[10]{0} %multiply.180), metadata={op_type=\"AssignAddVariableOp\" op_name=\"StatefulPartitionedCall_1/AssignAddVariableOp_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=195}\n", " %sqrt.182 = f32[10]{0} sqrt(f32[10]{0} %add.181), metadata={op_type=\"Sqrt\" op_name=\"Sqrt_1.StatefulPartitionedCall_1/Sqrt_1\"}\n", " %constant.184 = f32[] constant(1e-07), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall_1/add_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %broadcast.185 = f32[10]{0} broadcast(f32[] %constant.184), dimensions={}, metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall_1/add_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %add.186 = f32[10]{0} add(f32[10]{0} %sqrt.182, f32[10]{0} %broadcast.185), metadata={op_type=\"AddV2\" op_name=\"StatefulPartitionedCall_1/add_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %divide.187 = f32[10]{0} divide(f32[10]{0} %multiply.175, f32[10]{0} %add.186), metadata={op_type=\"RealDiv\" op_name=\"StatefulPartitionedCall_1/truediv_1\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %subtract.188 = f32[10]{0} subtract(f32[10]{0} %arg3.4, f32[10]{0} %divide.187), metadata={op_type=\"AssignSubVariableOp\" op_name=\"StatefulPartitionedCall_1/AssignSubVariableOp\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/adam.py\" source_line=200}\n", " %reshape.193 = f32[10]{0} reshape(f32[10]{0} %subtract.188), metadata={op_name=\"XLA_Retvals\"}\n", " %copy.194 = f32[10]{0} copy(f32[10]{0} %reshape.193), metadata={op_name=\"XLA_Retvals\"}\n", " %constant.189 = s64[] constant(1), metadata={op_type=\"AssignAddVariableOp\" op_name=\"AssignAddVariableOp\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/optimizer.py\" source_line=1236}\n", " %add.190 = s64[] add(s64[] %arg4.5, s64[] %constant.189), metadata={op_type=\"AssignAddVariableOp\" op_name=\"AssignAddVariableOp\" source_file=\"/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/keras/optimizers/optimizer_experimental/optimizer.py\" source_line=1236}\n", " %reshape.195 = s64[] reshape(s64[] %add.190), metadata={op_name=\"XLA_Retvals\"}\n", " %copy.196 = s64[] copy(s64[] %reshape.195), metadata={op_name=\"XLA_Retvals\"}\n", " %reshape.197 = f32[784,10]{1,0} reshape(f32[784,10]{1,0} %add.152), metadata={op_name=\"XLA_Retvals\"}\n", " %copy.198 = f32[784,10]{1,0} copy(f32[784,10]{1,0} %reshape.197), metadata={op_name=\"XLA_Retvals\"}\n", " %reshape.199 = f32[784,10]{1,0} reshape(f32[784,10]{1,0} %add.160), metadata={op_name=\"XLA_Retvals\"}\n", " %copy.200 = f32[784,10]{1,0} copy(f32[784,10]{1,0} %reshape.199), metadata={op_name=\"XLA_Retvals\"}\n", " %reshape.201 = f32[10]{0} reshape(f32[10]{0} %add.173), metadata={op_name=\"XLA_Retvals\"}\n", " %copy.202 = f32[10]{0} copy(f32[10]{0} %reshape.201), metadata={op_name=\"XLA_Retvals\"}\n", " %reshape.203 = f32[10]{0} reshape(f32[10]{0} %add.181), metadata={op_name=\"XLA_Retvals\"}\n", " %copy.204 = f32[10]{0} copy(f32[10]{0} %reshape.203), metadata={op_name=\"XLA_Retvals\"}\n", " ROOT %tuple.205 = (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} %copy.192, f32[10]{0} %copy.194, s64[] %copy.196, f32[784,10]{1,0} %copy.198, f32[784,10]{1,0} %copy.200, /*index=5*/f32[10]{0} %copy.202, f32[10]{0} %copy.204), metadata={op_name=\"XLA_Retvals\"}\n", "}\n", "\n", "\n" ] } ], "source": [ "print(train_mnist.experimental_get_compiler_ir(images, labels)(stage='hlo'))" ] } ], "metadata": { "colab": { "collapsed_sections": [], "name": "jit_compile.ipynb", "provenance": [], "toc_visible": true }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.9.16" } }, "nbformat": 4, "nbformat_minor": 0 }